Hhhh

Running

App Files Files Community

Hjgugugjhuhjggg commited on Mar 2

Commit

1c817fd

verified ·

1 Parent(s): bb7e407

Upload 27 files

Browse files

Files changed (27) hide show

Dockerfile +20 -0
README.md +14 -12
api.py +444 -0
background_tasks.py +197 -0
codegen_api.py +23 -0
configs.py +206 -0
constants.py +449 -0
extensions.py +252 -0
image_to_3d_api.py +32 -0
imagegen_api.py +33 -0
main.py +118 -0
model_loader.py +674 -0
models.py +96 -0
musicgen_api.py +35 -0
requirements.txt +40 -0
sadtalker_api.py +202 -0
sadtalker_utils.py +866 -0
sentiment_api.py +27 -0
stt_api.py +36 -0
summarization_api.py +29 -0
text_generation.py +152 -0
text_to_video_api.py +37 -0
tokenxxx.py +161 -0
translation_api.py +27 -0
tts_api.py +23 -0
utils.py +190 -0
xxx.py +142 -0

Dockerfile ADDED Viewed

	@@ -0,0 +1,20 @@

+FROM python:3.11-slim-buster
+ENV DEBIAN_FRONTEND=noninteractive
+ENV NUMBA_DISABLE_CACHE=1
+WORKDIR /app
+RUN apt-get update && apt-get upgrade -y
+RUN apt-get install libgl1-mesa-glx ffmpeg -y
+RUN mkdir -p /.cache/huggingface/hub && chmod -R 777 /.cache/huggingface/hub
+RUN mkdir -p /.config/matplotlib && chmod -R 777 /.config/matplotlib
+RUN mkdir -p /nltk_data && chmod -R 777 /nltk_data
+RUN pip install --no-cache-dir accelerate retry asyncio basicsr beautifulsoup4 bs4 opencv-python deep-translator duckduckgo-search fastapi flask flask-cors facexlib ffmpeg-python gfpgan imageio imageio-ffmpeg langdetect librosa nltk numpy Pillow pydub pytorch-lightning PyYAML retry safetensors scikit-learn scipy scikit-image soundfile torch torchaudio torchvision tqdm wget yacs numba
+COPY . .
+EXPOSE 7860
+CMD ["python", "main.py"]

README.md CHANGED Viewed

@@ -1,12 +1,14 @@
----
-title: Hhhh
-emoji: ⚡
-colorFrom: indigo
-colorTo: red
-sdk: docker
-pinned: false
-license: apache-2.0
-short_description: Apache2
----
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

+---
+title: Ggggggc
+emoji: 📈
+colorFrom: yellow
+colorTo: indigo
+sdk: docker
+sdk_version: 5.18.0
+app_file: main.py
+pinned: false
+license: apache-2.0
+short_description: Apache
+---
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

api.py ADDED Viewed

	@@ -0,0 +1,444 @@

+from main import *
+from tts_api import *
+from stt_api import *
+from sentiment_api import *
+from imagegen_api import *
+from musicgen_api import *
+from translation_api import *
+from codegen_api import *
+from text_to_video_api import *
+from summarization_api import *
+from image_to_3d_api import *
+from flask import Flask, request, jsonify, Response, send_file, stream_with_context
+from flask_cors import CORS
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+import numpy as np
+from PIL import Image
+import io
+import tempfile
+import queue
+import json
+import base64
+app = Flask(__name__)
+CORS(app)
+html_code = """<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>AI Text Generation</title>
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/animate.css/4.1.1/animate.min.css"/>
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha512-9usAa10IRO0HhonpyAIVpjrylPvoDwiPUiKdWk5t3PyolY1cOd4DSE0Ga+ri4AuTroPR5aQvXU9xC6qOPnzFeg==" crossorigin="anonymous" referrerpolicy="no-referrer" />
+    <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
+    <style>
+        body {
+            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+            background: #f0f0f0;
+            color: #333;
+            margin: 0;
+            padding: 0;
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+            min-height: 100vh;
+        }
+        .container {
+            width: 95%;
+            max-width: 900px;
+            padding: 20px;
+            background-color: #fff;
+            box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
+            border-radius: 8px;
+            margin-top: 20px;
+            margin-bottom: 20px;
+            display: flex;
+            flex-direction: column;
+        }
+        .header {
+            text-align: center;
+            margin-bottom: 20px;
+        }
+        .header h1 {
+            font-size: 2em;
+            color: #333;
+        }
+        .form-group {
+            margin-bottom: 15px;
+        }
+        .form-group textarea {
+            width: 100%;
+            padding: 10px;
+            border: 1px solid #ccc;
+            border-radius: 5px;
+            font-size: 16px;
+            box-sizing: border-box;
+            resize: vertical;
+        }
+        button {
+            padding: 10px 15px;
+            border: none;
+            border-radius: 5px;
+            background-color: #007bff;
+            color: white;
+            font-size: 18px;
+            cursor: pointer;
+            transition: background-color 0.3s ease;
+        }
+        button:hover {
+            background-color: #0056b3;
+        }
+        #output {
+            margin-top: 20px;
+            padding: 15px;
+            border: 1px solid #ddd;
+            border-radius: 5px;
+            background-color: #f9f9f9;
+            white-space: pre-wrap;
+            word-break: break-word;
+            overflow-y: auto;
+            max-height: 100vh;
+        }
+        #output strong {
+            font-weight: bold;
+        }
+        .animated-text {
+            position: fixed;
+            top: 20px;
+            left: 20px;
+            font-size: 1.5em;
+            color: rgba(0, 0, 0, 0.1);
+            pointer-events: none;
+            z-index: -1;
+        }
+        @media (max-width: 768px) {
+            .container {
+                width: 98%;
+                margin-top: 10px;
+                margin-bottom: 10px;
+                padding: 15px;
+            }
+            .header h1 {
+                font-size: 1.8em;
+            }
+            .form-group textarea, .form-group input[type="text"] {
+                font-size: 14px;
+                padding: 8px;
+            }
+            button {
+                font-size: 16px;
+                padding: 8px 12px;
+            }
+            #output {
+                font-size: 14px;
+                padding: 10px;
+                margin-top: 15px;
+            }
+        }
+    </style>
+</head>
+<body>
+<div class="animated-text animate__animated animate__fadeIn animate__infinite infinite">AI POWERED</div>
+<div class="container">
+    <div class="header animate__animated animate__fadeInDown">
+    </div>
+    <div class="form-group animate__animated animate__fadeInLeft">
+        <textarea id="text" rows="5" placeholder="Enter text"></textarea>
+    </div>
+    <button onclick="generateText()" class="animate__animated animate__fadeInUp">Generate Reasoning</button>
+    <div id="output" class="animate__animated">
+        <strong>Response:</strong><br>
+        <span id="generatedText"></span>
+    </div>
+</div>
+<script>
+    let eventSource = null;
+    let accumulatedText = "";
+    let lastResponse = "";
+    async function generateText() {
+        const inputText = document.getElementById("text").value;
+        document.getElementById("generatedText").innerText = "";
+        accumulatedText = "";
+        if (eventSource) {
+            eventSource.close();
+        }
+        const temp = 0.7;
+        const top_k_val = 40;
+        const top_p_val = 0.0;
+        const repetition_penalty_val = 1.2;
+        const requestData = {
+            text: inputText,
+            temp: temp,
+            top_k: top_k_val,
+            top_p: top_p_val,
+            reppenalty: repetition_penalty_val
+        };
+        const params = new URLSearchParams(requestData).toString();
+        eventSource = new EventSource('/api/v1/generate_stream?' + params);
+        eventSource.onmessage = function(event) {
+            if (event.data === "<END_STREAM>") {
+                eventSource.close();
+                const currentResponse = accumulatedText.replace("<|endoftext|>", "").replace(/\s+(?=[.,，。])/g, '').trim();
+                if (currentResponse === lastResponse.trim()) {
+                    accumulatedText = "**Response is repetitive. Please try again or rephrase your query.**";
+                } else {
+                    lastResponse = currentResponse;
+                }
+                document.getElementById("generatedText").innerHTML = marked.parse(accumulatedText);
+                return;
+            }
+            accumulatedText += event.data;
+            let partialText = accumulatedText.replace("<|endoftext|>", "").replace(/\s+(?=[.,，。])/g, '').trim();
+            document.getElementById("generatedText").innerHTML = marked.parse(partialText);
+        };
+        eventSource.onerror = function(error) {
+            console.error("SSE error", error);
+            eventSource.close();
+        };
+        const outputDiv = document.getElementById("output");
+        outputDiv.classList.add("show");
+    }
+    function base64ToBlob(base64Data, contentType) {
+        contentType = contentType || '';
+        const sliceSize = 1024;
+        const byteCharacters = atob(base64Data);
+        const bytesLength = byteCharacters.length;
+        const slicesCount = Math.ceil(bytesLength / sliceSize);
+        const byteArrays = new Array(slicesCount);
+        for (let sliceIndex = 0; sliceIndex < slicesCount; ++sliceIndex) {
+            const begin = sliceIndex * sliceSize;
+            const end = Math.min(begin + sliceSize, bytesLength);
+            const bytes = new Array(end - begin);
+            for (let offset = begin, i = 0; offset < end; ++i, ++offset) {
+                bytes[i] = byteCharacters[offset].charCodeAt(0);
+            }
+            byteArrays[sliceIndex] = new Uint8Array(bytes);
+        }
+        return new Blob(byteArrays, { type: contentType });
+    }
+</script>
+</body>
+</html>
+"""
+feedback_queue = queue.Queue()
+class TextGenerationModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim, hidden_dim):
+        super(TextGenerationModel, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.rnn = nn.GRU(embed_dim, hidden_dim, batch_first=True)
+        self.fc = nn.Linear(hidden_dim, vocab_size)
+    def forward(self, x, hidden=None):
+        x = self.embedding(x)
+        out, hidden = self.rnn(x, hidden)
+        out = self.fc(out)
+        return out, hidden
+vocab = ["hola", "mundo", "este", "es", "un", "ejemplo", "de", "texto", "generado", "con", "torch"]
+vocab_size = len(vocab)
+embed_dim = 16
+hidden_dim = 32
+text_model = TextGenerationModel(vocab_size, embed_dim, hidden_dim)
+text_model.eval()
+def tokenize(text):
+    tokens = text.lower().split()
+    indices = [vocab.index(token) if token in vocab else 0 for token in tokens]
+    return torch.tensor(indices, dtype=torch.long).unsqueeze(0)
+def perform_reasoning_stream(text, temperature, top_k, top_p, repetition_penalty):
+    input_tensor = tokenize(text)
+    hidden = None
+    for _ in range(20):
+        outputs, hidden = text_model(input_tensor, hidden)
+        logits = outputs[:, -1, :] / temperature
+        probs = F.softmax(logits, dim=-1)
+        topk_probs, topk_indices = torch.topk(probs, min(top_k, logits.shape[-1]))
+        chosen_index = topk_indices[0, torch.multinomial(topk_probs[0], 1).item()].item()
+        token_str = vocab[chosen_index]
+        yield token_str
+        input_tensor = torch.cat([input_tensor, torch.tensor([[chosen_index]], dtype=torch.long)], dim=1)
+    yield "<END_STREAM>"
+class SentimentModel(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim):
+        super(SentimentModel, self).__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, output_dim)
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+sentiment_model = SentimentModel(10, 16, 2)
+sentiment_model.eval()
+@app.route("/")
+def index():
+    return html_code
+@app.route("/api/v1/generate_stream", methods=["GET"])
+def generate_stream():
+    text = request.args.get("text", "")
+    temp = float(request.args.get("temp", 0.7))
+    top_k = int(request.args.get("top_k", 40))
+    top_p = float(request.args.get("top_p", 0.0))
+    reppenalty = float(request.args.get("reppenalty", 1.2))
+    @stream_with_context
+    def event_stream():
+        try:
+            for token in perform_reasoning_stream(text, temperature=temp, top_k=top_k, top_p=top_p, repetition_penalty=reppenalty):
+                if token == "<END_STREAM>":
+                    yield "data: <END_STREAM>\n\n"
+                    break
+                yield "data: " + token + "\n\n"
+        except Exception as e:
+            yield "data: <ERROR>\n\n"
+    return Response(event_stream(), mimetype="text/event-stream")
+@app.route("/api/v1/generate", methods=["POST"])
+def generate():
+    data = request.get_json()
+    text = data.get("text", "")
+    temp = float(data.get("temp", 0.7))
+    top_k = int(data.get("top_k", 40))
+    top_p = float(data.get("top_p", 0.0))
+    reppenalty = float(data.get("reppenalty", 1.2))
+    result = ""
+    try:
+        for token in perform_reasoning_stream(text, temperature=temp, top_k=top_k, top_p=top_p, repetition_penalty=reppenalty):
+            if token == "<END_STREAM>":
+                break
+            result += token + " "
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+    return jsonify({"solidity": result.strip()})
+@app.route("/api/v1/feedback", methods=["POST"])
+def feedback():
+    data = request.get_json()
+    feedback_text = data.get("feedback_text")
+    correct_category = data.get("correct_category")
+    if feedback_text and correct_category:
+        feedback_queue.put((feedback_text, correct_category))
+        return jsonify({"status": "feedback received"})
+    return jsonify({"status": "feedback failed"}), 400
+@app.route("/api/v1/tts", methods=["POST"])
+def tts_api():
+    data = request.get_json()
+    text = data.get("text", "")
+    sr = 22050
+    duration = 3.0
+    t = torch.linspace(0, duration, int(sr * duration))
+    frequency = 440.0
+    audio = 0.5 * torch.sin(2 * torch.pi * frequency * t)
+    audio = audio.unsqueeze(0)
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+        torchaudio.save(tmp.name, audio, sr)
+        tmp_path = tmp.name
+    return send_file(tmp_path, mimetype="audio/wav", as_attachment=True, download_name="output.wav")
+@app.route("/api/v1/stt", methods=["POST"])
+def stt_api():
+    data = request.get_json()
+    audio_b64 = data.get("audio", "")
+    if audio_b64:
+        audio_bytes = base64.b64decode(audio_b64)
+        buf = io.BytesIO(audio_bytes)
+        waveform, sr = torchaudio.load(buf)
+        mean_amp = waveform.abs().mean().item()
+        recognized_text = f"Audio processed with mean amplitude {mean_amp:.3f}"
+        return jsonify({"text": recognized_text})
+    return jsonify({"text": ""})
+@app.route("/api/v1/sentiment", methods=["POST"])
+def sentiment_api():
+    data = request.get_json()
+    text = data.get("text", "")
+    if not text:
+        return jsonify({"sentiment": "neutral"})
+    ascii_vals = [ord(c) for c in text[:10]]
+    while len(ascii_vals) < 10:
+        ascii_vals.append(0)
+    features = torch.tensor(ascii_vals, dtype=torch.float32).unsqueeze(0)
+    output = sentiment_model(features)
+    sentiment_idx = torch.argmax(output, dim=1).item()
+    sentiment = "positivo" if sentiment_idx == 1 else "negativo"
+    return jsonify({"sentiment": sentiment})
+@app.route("/api/v1/imagegen", methods=["POST"])
+def imagegen_api():
+    data = request.get_json()
+    prompt = data.get("prompt", "")
+    image_tensor = torch.rand(3, 256, 256)
+    np_image = image_tensor.mul(255).clamp(0, 255).byte().numpy().transpose(1, 2, 0)
+    img = Image.fromarray(np_image)
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    buf.seek(0)
+    return send_file(buf, mimetype="image/png", as_attachment=True, download_name="image.png")
+@app.route("/api/v1/musicgen", methods=["POST"])
+def musicgen_api():
+    data = request.get_json()
+    prompt = data.get("prompt", "")
+    sr = 22050
+    duration = 5.0
+    t = torch.linspace(0, duration, int(sr * duration))
+    frequency = 440.0
+    audio = 0.5 * torch.sin(2 * torch.pi * frequency * t)
+    audio = audio.unsqueeze(0)
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+        torchaudio.save(tmp.name, audio, sr)
+        tmp_path = tmp.name
+    return send_file(tmp_path, mimetype="audio/wav", as_attachment=True, download_name="music.wav")
+@app.route("/api/v1/translation", methods=["POST"])
+def translation_api():
+    data = request.get_json()
+    text = data.get("text", "")
+    translated = " ".join(text.split()[::-1])
+    return jsonify({"translated_text": translated})
+@app.route("/api/v1/codegen", methods=["POST"])
+def codegen_api():
+    data = request.get_json()
+    prompt = data.get("prompt", "")
+    generated_code = f"# Generated code based on prompt: {prompt}\nprint('Hello from Torch-generated code')"
+    return jsonify({"code": generated_code})
+@app.route("/api/v1/text_to_video", methods=["POST"])
+def text_to_video_api():
+    data = request.get_json()
+    prompt = data.get("prompt", "")
+    video_tensor = torch.randint(0, 255, (10, 3, 64, 64), dtype=torch.uint8)
+    video_bytes = video_tensor.numpy().tobytes()
+    buf = io.BytesIO(video_bytes)
+    return send_file(buf, mimetype="video/mp4", as_attachment=True, download_name="video.mp4")
+@app.route("/api/v1/summarization", methods=["POST"])
+def summarization_api():
+    data = request.get_json()
+    text = data.get("text", "")
+    sentences = text.split('.')
+    summary = sentences[0] if sentences[0] else text
+    return jsonify({"summary": summary})
+@app.route("/api/v1/image_to_3d", methods=["POST"])
+def image_to_3d_api():
+    data = request.get_json()
+    prompt = data.get("prompt", "")
+    obj_data = "o Cube\nv 0 0 0\nv 1 0 0\nv 1 1 0\nv 0 1 0\nf 1 2 3 4"
+    buf = io.BytesIO(obj_data.encode("utf-8"))
+    return send_file(buf, mimetype="text/plain", as_attachment=True, download_name="model.obj")
+@app.route("/api/v1/sadtalker", methods=["GET"])
+def sadtalker():
+    return jsonify({"message": "Respuesta de sadtalker"})
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)

background_tasks.py ADDED Viewed

	@@ -0,0 +1,197 @@

+import time
+import threading
+import queue
+import uuid
+import unicodedata
+import re
+from deep_translator import GoogleTranslator
+from duckduckgo_search import DDGS
+import nltk
+import torch
+import torch.nn as nn
+import math
+nltk.download('punkt')
+categories = ['News', 'Sports', 'Entertainment']
+TEXT_GENERATION_RATE = 10
+text_queue = queue.Queue()
+reasoning_queue = queue.Queue()
+feedback_queue = queue.Queue()
+vocabulary = ["<PAD>", "<EOS>"]
+word_to_index = {word: idx for idx, word in enumerate(vocabulary)}
+seen_responses = set()
+news_clf = None
+class SimpleClassifier(nn.Module):
+    def __init__(self, vocab_size, num_classes, embedding_dim=128):
+        super(SimpleClassifier, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embedding_dim)
+        self.fc = nn.Linear(embedding_dim, num_classes)
+    def forward(self, x):
+        embedded = self.embedding(x)
+        pooled = embedded.mean(dim=1)
+        out = self.fc(pooled)
+        return out
+def tokenize_text(text):
+    return nltk.word_tokenize(text)
+def update_vocabulary(tokens):
+    global vocabulary, word_to_index
+    for token in tokens:
+        if token not in word_to_index:
+            word_to_index[token] = len(vocabulary)
+            vocabulary.append(token)
+def text_to_vector(text):
+    tokens = tokenize_text(text)
+    update_vocabulary(tokens)
+    indices = [word_to_index.get(token, 0) for token in tokens]
+    return torch.tensor(indices, dtype=torch.long)
+def generate_and_queue_text(language):
+    global categories, text_queue
+    num_categories = len(categories)
+    num_texts_per_category = TEXT_GENERATION_RATE // (2 * num_categories)
+    while True:
+        for category in categories:
+            for _ in range(num_texts_per_category):
+                uid = uuid.uuid4()
+                base_text = f"Category: {category}. ID:{uid}"
+                try:
+                    translator = GoogleTranslator(source='auto', target=language)
+                    text = translator.translate(base_text)
+                except Exception:
+                    text = base_text
+                processed_text = ''.join(c for c in unicodedata.normalize('NFKC', text) if c.isprintable())
+                text_queue.put((processed_text, category))
+                time.sleep(0)
+def background_training():
+    global categories, news_clf, feedback_queue, vocabulary
+    if categories is None:
+        categories = ['DefaultCategory']
+    num_classes = len(categories)
+    learning_rate = 0.01
+    epochs = 1
+    if news_clf is None:
+        news_clf = SimpleClassifier(len(vocabulary), num_classes)
+    optimizer = torch.optim.SGD(news_clf.parameters(), lr=learning_rate)
+    criterion = nn.CrossEntropyLoss()
+    while True:
+        try:
+            feedback_item = feedback_queue.get(timeout=10)
+            if feedback_item:
+                input_text, generated_text = feedback_item
+                input_vector = text_to_vector(input_text)
+                if len(vocabulary) == 0:
+                    vocabulary.extend(["<PAD>", "<EOS>"])
+                    news_clf = SimpleClassifier(len(vocabulary), num_classes)
+                    optimizer = torch.optim.SGD(news_clf.parameters(), lr=learning_rate)
+                if input_vector.size(0) != len(vocabulary) and len(vocabulary) > 0:
+                    news_clf = SimpleClassifier(len(vocabulary), num_classes)
+                    optimizer = torch.optim.SGD(news_clf.parameters(), lr=learning_rate)
+                    input_vector = text_to_vector(input_text)
+                tokens = tokenize_text(input_text)
+                update_vocabulary(tokens)
+                tokens_indices = [word_to_index.get(word, 0) for word in tokens]
+                input_tensor = torch.tensor([tokens_indices], dtype=torch.long)
+                target_index = categories.index(generated_text) if generated_text in categories else 0
+                target_category_index = torch.tensor([target_index], dtype=torch.long)
+                if num_classes <= 1:
+                    num_classes = 2
+                    news_clf.fc = nn.Linear(128, num_classes)
+                for _ in range(epochs):
+                    optimizer.zero_grad()
+                    output = news_clf(input_tensor)
+                    loss = criterion(output, target_category_index)
+                    loss.backward()
+                    optimizer.step()
+                feedback_queue.task_done()
+        except queue.Empty:
+            pass
+        except Exception:
+            time.sleep(5)
+class ReasoningModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim=128, hidden_dim=128):
+        super(ReasoningModel, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.rnn = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
+        self.fc = nn.Linear(hidden_dim, vocab_size)
+    def forward(self, x, hidden=None):
+        emb = self.embedding(x)
+        output, hidden = self.rnn(emb, hidden)
+        logits = self.fc(output)
+        return logits, hidden
+    def generate(self, input_seq, max_length=50, temperature=1.0):
+        self.eval()
+        tokens = input_seq.copy()
+        hidden = None
+        generated = []
+        for _ in range(max_length):
+            input_tensor = torch.tensor([tokens], dtype=torch.long)
+            logits, hidden = self.forward(input_tensor, hidden)
+            next_token_logits = logits[0, -1, :] / temperature
+            probabilities = torch.softmax(next_token_logits, dim=0)
+            next_token = torch.multinomial(probabilities, 1).item()
+            tokens.append(next_token)
+            generated.append(next_token)
+            if next_token == word_to_index.get("<EOS>"):
+                break
+        return generated
+reasoning_model = ReasoningModel(len(vocabulary))
+def perform_reasoning_stream(text_input, temperature=0.7, top_k=40, top_p=0.0, repetition_penalty=1.2):
+    tokens = tokenize_text(text_input)
+    update_vocabulary(tokens)
+    tokens_indices = [word_to_index.get(token, 0) for token in tokens]
+    generated_indices = reasoning_model.generate(tokens_indices, max_length=50, temperature=temperature)
+    for idx in generated_indices:
+        yield vocabulary[idx] + " "
+    yield "<END_STREAM>"
+def background_reasoning_queue():
+    global reasoning_queue, seen_responses
+    while True:
+        try:
+            item = reasoning_queue.get(timeout=1)
+            if item is None:
+                reasoning_queue.task_done()
+                continue
+            text_input = item.get('text_input')
+            temperature = item.get('temperature', 0.7)
+            top_k = item.get('top_k', 40)
+            top_p = item.get('top_p', 0.0)
+            repetition_penalty = item.get('repetition_penalty', 1.2)
+            resp_queue = item.get('response_queue', queue.Queue())
+            if not text_input:
+                resp_queue.put({"error": "Empty text input received."})
+                reasoning_queue.task_done()
+                continue
+            generated_text_stream = perform_reasoning_stream(text_input, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty)
+            full_response = ""
+            for chunk in generated_text_stream:
+                if chunk == "<END_STREAM>":
+                    break
+                full_response += chunk
+            cleaned_response = re.sub(r'\s+(?=[.,，。])', '', full_response.replace("<|endoftext|>", "")).strip()
+            if cleaned_response in seen_responses:
+                final_response = "**Response is repetitive. Please try again or rephrase your query.**"
+                resp_queue.put({"text": final_response})
+            else:
+                seen_responses.add(cleaned_response)
+                final_response = cleaned_response
+                resp_queue.put({"text": final_response})
+            reasoning_queue.task_done()
+        except queue.Empty:
+            pass
+        except Exception as e:
+            try:
+                resp_queue.put({"error": str(e)})
+            except Exception:
+                pass
+            if reasoning_queue and not reasoning_queue.empty():
+                reasoning_queue.task_done()

codegen_api.py ADDED Viewed

	@@ -0,0 +1,23 @@

+from flask import jsonify, send_file, request
+from main import *
+#from main import import codegen_model, codegen_tokenizer, device
+def generate_code(prompt, output_path="output_code.py"):
+    if codegen_model is None:
+        return "Code generation model not initialized."
+    input_ids = codegen_tokenizer.encode(prompt, return_tensors='pt').to(device)
+    output = codegen_model.generate(input_ids, max_length=512, temperature=0.7, top_p=0.9)
+    code = codegen_tokenizer.decode(output[0], skip_special_tokens=True)
+    with open(output_path, "w") as file:
+        file.write(code)
+    return output_path
+def codegen_api():
+    data = request.get_json()
+    prompt = data.get('prompt')
+    if not prompt:
+        return jsonify({"error": "Prompt is required"}), 400
+    output_file = generate_code(prompt)
+    if output_file == "Code generation model not initialized.":
+        return jsonify({"error": "Code generation failed"}), 500
+    return send_file(output_file, mimetype="text/x-python", as_attachment=True, download_name="output.py")

configs.py ADDED Viewed

	@@ -0,0 +1,206 @@

+from constants import *
+class GPT2Config:
+    def __init__(self, vocab_size_or_config_json_file=50257, n_positions=MAX_LENGTH, n_ctx=MAX_LENGTH, n_embd=768, n_layer=12, n_head=12, layer_norm_epsilon=1e-05, initializer_range=0.02):
+        self.vocab_size = vocab_size_or_config_json_file
+        self.n_ctx = n_ctx
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+    @classmethod
+    def from_dict(cls, config_dict):
+        return cls(**config_dict)
+class MBartConfig:
+    def __init__(self, vocab_size, d_model, num_layers, num_heads, pad_token_id, eos_token_id):
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.encoder_layers = num_layers
+        self.decoder_layers = num_layers
+        self.encoder_attention_heads = num_heads
+        self.decoder_attention_heads = num_heads
+        self.encoder_ffn_dim = d_model * 4
+        self.decoder_ffn_dim = d_model * 4
+        self.dropout = 0.1
+        self.attention_dropout = 0.0
+        self.activation_dropout = 0.0
+        self.max_position_embeddings = 1024
+        self.init_std = 0.02
+        self.layer_norm_eps = 1e-5
+        self.pad_token_id = pad_token_id
+        self.eos_token_id = eos_token_id
+        self.bos_token_id = 0
+        self.decoder_start_token_id = 2
+        self.output_past = True
+        self.scale_embedding = True
+        self.use_cache = True
+        self.num_hidden_layers = num_layers
+class CodeGenConfig:
+    def __init__(self, vocab_size, n_embd, n_layer, n_head):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_positions = 2048
+        self.resid_pdrop = 0.1
+        self.embd_pdrop = 0.1
+        self.attn_pdrop = 0.1
+        self.activation_function = "gelu_new"
+        self.n_ctx = 2048
+        self.pad_token_id = 50256
+        self.eos_token_id = 50256
+        self.initializer_range = 0.02
+class SummarizationConfig:
+    def __init__(self):
+        self.vocab_size = 10000
+        self.embedding_dim = 256
+        self.hidden_dim = 512
+        self.num_layers = 2
+        self.max_seq_len = 512
+class Clip4ClipConfig:
+    def __init__(self, vocab_size=30522, hidden_size=512, num_hidden_layers=6, num_attention_heads=8, intermediate_size=2048, hidden_act="gelu", hidden_dropout_prob=0.0, attention_probs_dropout_prob=0.0, max_position_embeddings=512, type_vocab_size=2, initializer_range=0.02, layer_norm_eps=1e-12, pad_token_id=0, bos_token_id=1, eos_token_id=2, **kwargs):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.all_head_size = self.num_attention_heads * self.hidden_size
+        self.attention_head_size = int(self.hidden_size / self.num_attention_heads)
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+    @classmethod
+    def from_dict(cls, config_dict):
+        return cls(**config_dict)
+class MusicGenConfig:
+    def __init__(self, vocab_size=2048, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act="gelu", hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, layer_norm_eps=1e-05, initializer_range=0.02, pad_token_id=0, bos_token_id=1, eos_token_id=2, n_positions=2048, n_ctx=2048, **kwargs):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.layer_norm_eps = layer_norm_eps
+        self.initializer_range = initializer_range
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.n_positions = n_positions
+        self.n_ctx = n_ctx
+        self.all_head_size = self.num_attention_heads * self.hidden_size
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+    @classmethod
+    def from_dict(cls, config_dict):
+        return cls(**config_dict)
+class BartConfig:
+    def __init__(self, vocab_size=50265, max_position_embeddings=1024, encoder_layers=12, encoder_ffn_dim=4096, encoder_attention_heads=16, decoder_layers=12, decoder_ffn_dim=4096, decoder_attention_heads=16, encoder_layerdrop=0.0, decoder_layerdrop=0.0, activation_function="gelu", d_model=1024, dropout=0.1, attention_dropout=0.0, activation_dropout=0.0, init_std=0.02, classifier_dropout=0.0, num_labels=3, pad_token_id=1, bos_token_id=0, eos_token_id=2, layer_norm_eps=1e-05, num_beams=4, early_stopping=True, max_length=100, min_length=30, scale_embedding=False, **kwargs):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.encoder_layers = encoder_layers
+        self.encoder_ffn_dim = encoder_ffn_dim
+        self.encoder_attention_heads = encoder_attention_heads
+        self.decoder_layers = decoder_layers
+        self.decoder_ffn_dim = decoder_ffn_dim
+        self.decoder_attention_heads = decoder_attention_heads
+        self.encoder_layerdrop = encoder_layerdrop
+        self.decoder_layerdrop = decoder_layerdrop
+        self.activation_function = activation_function
+        self.d_model = d_model
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.init_std = init_std
+        self.classifier_dropout = classifier_dropout
+        self.num_labels = num_labels
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.layer_norm_eps = layer_norm_eps
+        self.num_beams = num_beams
+        self.early_stopping = True
+        self.max_length = max_length
+        self.min_length = min_length
+        self.scale_embedding = False
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+    @classmethod
+    def from_dict(cls, config_dict):
+        return cls(**config_dict)
+class OpenLRMConfig:
+    def __init__(self, obj_dim=1024, hidden_dim=512, num_layers=6, num_heads=8, dropout_prob=0.1, **kwargs):
+        self.obj_dim = obj_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.dropout_prob = dropout_prob
+        self.all_head_size = self.num_heads * self.hidden_dim
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+    @classmethod
+    def from_dict(cls, config_dict):
+        return cls(**config_dict)
+class UNet2DConditionModelConfig:
+    def __init__(self, sample_size=64, layers_per_block=2, block_out_channels=[320, 640, 1280, 1280], downsample=[2, 2, 2, 2], upsample=[2, 2, 2, 2], cross_attention_dim=768, act_fn="silu", norm_num_groups=32, num_attention_heads=8, in_channels=4, out_channels=4, attention_head_dim=64, **kwargs):
+        self.sample_size = sample_size
+        self.layers_per_block = layers_per_block
+        self.block_out_channels = block_out_channels
+        self.downsample = downsample
+        self.upsample = upsample
+        self.cross_attention_dim = cross_attention_dim
+        self.act_fn = act_fn
+        self.norm_num_groups = norm_num_groups
+        self.num_attention_heads = num_attention_heads
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.attention_head_dim = attention_head_dim
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+    @classmethod
+    def from_dict(cls, config_dict):
+        return cls(**config_dict)
+class AutoencoderKLConfig:
+    def __init__(self, **kwargs):
+        self.sample_size = 64
+        self.latent_channels = 4
+        self.layers_per_block = 2
+        self.block_out_channels = [128, 256, 512, 512]
+        self.downsample = [2, 2, 2, 2]
+        self.upsample = [2, 2, 2, 2]
+        self.act_fn = "silu"
+        self.norm_num_groups = 32
+        self.num_channels_every_n_layers = 2
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+    @classmethod
+    def from_dict(cls, config_dict):
+        return cls(**config_dict)

constants.py ADDED Viewed

	@@ -0,0 +1,449 @@

+import os
+TEXT_GENERATION_RATE = 40000
+MAX_LENGTH = 2048
+MAX_XDD = 5
+END_OF_TEXT_TOKEN = "<|endoftext|>"
+SYSTEM_PROMPT = """Eres un asistente experto con habilidades avanzadas en diversas áreas. Responde de manera amigable, educada y razonada. Siempre piensa cuidadosamente antes de responder para asegurar la claridad y completitud. Posees la capacidad de autoaprendizaje continuo y recuerdas interacciones pasadas para mejorar tus respuestas y evitar errores repetidos."""
+XML_COT_FORMAT = """<reasoning>\n{reasoning}\n</reasoning>\n<answer>\n{answer}\n</answer>\n"""
+html_code = """<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>AI Text Generation</title>
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/animate.css/4.1.1/animate.min.css"/>
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha512-9usAa10IRO0HhonpyAIVpjrylPvoDwiPUiKdWk5t3PyolY1cOd4DSE0Ga+ri4AuTroPR5aQvXU9xC6qOPnzFeg==" crossorigin="anonymous" referrerpolicy="no-referrer" />
+    <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
+    <style>
+        body {
+            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+            background: #f0f0f0;
+            color: #333;
+            margin: 0;
+            padding: 0;
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+            min-height: 100vh;
+        }
+        .container {
+            width: 95%;
+            max-width: 900px;
+            padding: 20px;
+            background-color: #fff;
+            box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
+            border-radius: 8px;
+            margin-top: 20px;
+            margin-bottom: 20px;
+            display: flex;
+            flex-direction: column;
+        }
+        .header {
+            text-align: center;
+            margin-bottom: 20px;
+        }
+        .header h1 {
+            font-size: 2em;
+            color: #333;
+        }
+        .form-group {
+            margin-bottom: 15px;
+        }
+        .form-group textarea {
+            width: 100%;
+            padding: 10px;
+            border: 1px solid #ccc;
+            border-radius: 5px;
+            font-size: 16px;
+            box-sizing: border-box;
+            resize: vertical;
+        }
+        button {
+            padding: 10px 15px;
+            border: none;
+            border-radius: 5px;
+            background-color: #007bff;
+            color: white;
+            font-size: 18px;
+            cursor: pointer;
+            transition: background-color 0.3s ease;
+        }
+        button:hover {
+            background-color: #0056b3;
+        }
+        #output {
+            margin-top: 20px;
+            padding: 15px;
+            border: 1px solid #ddd;
+            border-radius: 5px;
+            background-color: #f9f9f9;
+            white-space: pre-wrap;
+            word-break: break-word;
+            overflow-y: auto;
+            max-height: 100vh;
+        }
+        #output strong {
+            font-weight: bold;
+        }
+        .animated-text {
+            position: fixed;
+            top: 20px;
+            left: 20px;
+            font-size: 1.5em;
+            color: rgba(0, 0, 0, 0.1);
+            pointer-events: none;
+            z-index: -1;
+        }
+        @media (max-width: 768px) {
+            .container {
+                width: 98%;
+                margin-top: 10px;
+                margin-bottom: 10px;
+                padding: 15px;
+            }
+            .header h1 {
+                font-size: 1.8em;
+            }
+            .form-group textarea, .form-group input[type="text"] {
+                font-size: 14px;
+                padding: 8px;
+            }
+            button {
+                font-size: 16px;
+                padding: 8px 12px;
+            }
+            #output {
+                font-size: 14px;
+                padding: 10px;
+                margin-top: 15px;
+            }
+        }
+    </style>
+</head>
+<body>
+<div class="animated-text animate__animated animate__fadeIn animate__infinite infinite">AI POWERED</div>
+<div class="container">
+    <div class="header animate__animated animate__fadeInDown">
+    </div>
+    <div class="form-group animate__animated animate__fadeInLeft">
+        <textarea id="text" rows="5" placeholder="Enter text"></textarea>
+    </div>
+    <button onclick="generateText()" class="animate__animated animate__fadeInUp">Generate Reasoning</button>
+    <div id="output" class="animate__animated">
+        <strong >Response:</strong><br>
+        <span id="generatedText"></span>
+    </div>
+</div>
+<script>
+    let eventSource = null;
+    let accumulatedText = "";
+    let lastResponse = "";
+    async function generateText() {
+        const inputText = document.getElementById("text").value;
+        document.getElementById("generatedText").innerText = "";
+        accumulatedText = "";
+        if (eventSource) {
+            eventSource.close();
+        }
+        const temp = 0.7;
+        const top_k_val = 40;
+        const top_p_val = 0.0;
+        const repetition_penalty_val = 1.2;
+        const requestData = {
+            text: inputText,
+            temp: temp,
+            top_k: top_k_val,
+            top_p: top_p_val,
+            reppenalty: repetition_penalty_val
+        };
+        eventSource = new EventSource('/generate_stream', {
+            headers: {
+                'Content-Type': 'application/json'
+            },
+            method: 'POST',
+            body: JSON.stringify(requestData)
+        });
+        eventSource.onmessage = function(event) {
+            if (event.data === "<END_STREAM>") {
+                eventSource.close();
+                const currentResponse = accumulatedText.replace("<|endoftext|>", "").replace(re.compile(r'\\s+(?=[.,，。])'), '').trim();
+                if (currentResponse === lastResponse.trim()) {
+                    accumulatedText = "**Response is repetitive. Please try again or rephrase your query.**";
+                } else {
+                    lastResponse = currentResponse;
+                }
+                document.getElementById("generatedText").innerHTML = marked.parse(accumulatedText);
+                return;
+            }
+            accumulatedText += event.data;
+            let partialText = accumulatedText.replace("<|endoftext|>", "").replace(re.compile(r'\\s+(?=[.,，。])'), '').trim();
+            document.getElementById("generatedText").innerHTML = marked.parse(partialText);
+        };
+        eventSource.onerror = function(error) {
+            console.error("SSE error", error);
+            eventSource.close();
+        };
+        const outputDiv = document.getElementById("output");
+        outputDiv.classList.add("show");
+    }
+    function base64ToBlob(base64Data, contentType) {
+        contentType = contentType || '';
+        const sliceSize = 1024;
+        const byteCharacters = atob(base64Data);
+        const bytesLength = byteCharacters.length;
+        const slicesCount = Math.ceil(bytesLength / sliceSize);
+        const byteArrays = new Array(slicesCount);
+        for (let sliceIndex = 0; sliceIndex < slicesCount; ++sliceIndex) {
+            const begin = sliceIndex * sliceSize;
+            const end = Math.min(begin + sliceSize, bytesLength);
+            const bytes = new Array(end - begin);
+            for (let offset = begin, i = 0; offset < end; ++i, ++offset) {
+                bytes[i] = byteCharacters[offset].charCodeAt(0);
+            }
+            byteArrays[sliceIndex] = new Uint8Array(bytes);
+        }
+        return new Blob(byteArrays, { type: contentType });
+    }
+</script>
+</body>
+</html>
+"""
+HTML_CODE = html_code
+# =============================================================================
+# Constantes definidas por el usuario
+# =============================================================================
+# GPT-2
+GPT2_FOLDER = "./GPT2"
+MODEL_FILE = "gpt2-pytorch_model.bin"
+ENCODER_FILE = "encoder.json"
+VOCAB_FILE = "vocab.bpe"
+CONFIG_FILE = "config.json"
+GPT2CONFHG = "https://huggingface.co/openai-community/gpt2/resolve/main/config.json"
+MODEL_URL = "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-pytorch_model.bin"
+ENCODER_URL = "https://raw.githubusercontent.com/graykode/gpt-2-Pytorch/refs/heads/master/GPT2/encoder.json"
+VOCAB_URL = "https://raw.githubusercontent.com/graykode/gpt-2-Pytorch/refs/heads/master/GPT2/vocab.bpe"
+# Traducción (MBart)
+TRANSLATION_FOLDER = "./TranslationModel"
+TRANSLATION_MODEL_WEIGHTS_FILE = "pytorch_model.bin"
+TRANSLATION_MODEL_CONFIG_FILE = "config.json"
+TRANSLATION_MODEL_VOCAB_FILE = "sentencepiece.bpe.model"
+TRANSLATION_MODEL_WEIGHTS_URL = "https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt/resolve/main/pytorch_model.bin"
+TRANSLATION_MODEL_CONFIG_URL = "https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt/resolve/main/config.json"
+TRANSLATION_MODEL_VOCAB_URL = "https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt/resolve/main/sentencepiece.bpe.model"
+TRANSLATION_MODEL_FILES_URLS = [
+    (TRANSLATION_MODEL_WEIGHTS_URL, TRANSLATION_MODEL_WEIGHTS_FILE),
+    (TRANSLATION_MODEL_CONFIG_URL, TRANSLATION_MODEL_CONFIG_FILE),
+    (TRANSLATION_MODEL_VOCAB_URL, TRANSLATION_MODEL_VOCAB_FILE),
+]
+# CodeGen
+CODEGEN_FOLDER = "./CodeGenModel"
+CODEGEN_MODEL_NAME = "codegen-350M-multi"
+CODEGEN_MODEL_WEIGHTS = "pytorch_model.bin"
+CODEGEN_CONFIG = "config.json"
+CODEGEN_VOCAB = "vocab.json"
+CODEGEN_MERGES = "merges.txt"
+CODEGEN_MODEL_WEIGHTS_URL = "https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/pytorch_model.bin"
+CODEGEN_CONFIG_URL = "https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json"
+CODEGEN_VOCAB_URL = "https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/vocab.json"
+CODEGEN_MERGES_URL = "https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/merges.txt"
+CODEGEN_FILES_URLS = [
+    (CODEGEN_MODEL_WEIGHTS_URL, CODEGEN_MODEL_WEIGHTS),
+    (CODEGEN_CONFIG_URL, CODEGEN_CONFIG),
+    (CODEGEN_VOCAB_URL, CODEGEN_VOCAB),
+    (CODEGEN_MERGES_URL, CODEGEN_MERGES),
+]
+# MusicGen
+MUSICGEN_FOLDER = "./MusicGenModel"
+MUSICGEN_MODEL_NAME = "melody"
+MUSICGEN_MODEL_WEIGHTS = "pytorch_model.bin"
+MUSICGEN_CONFIG = "config.json"
+MUSICGEN_SAMPLE_RATE = 32000
+MUSICGEN_DURATION = 8
+MUSICGEN_MODEL_WEIGHTS_URL = "https://huggingface.co/facebook/musicgen-small/resolve/main/pytorch_model.bin"
+MUSICGEN_CONFIG_URL = "https://huggingface.co/facebook/musicgen-small/resolve/main/config.json"
+MUSICGEN_FILES_URLS = [
+    (MUSICGEN_MODEL_WEIGHTS_URL, MUSICGEN_MODEL_WEIGHTS),
+    (MUSICGEN_CONFIG_URL, MUSICGEN_CONFIG)
+]
+# Summarization (Bart)
+SUMMARIZATION_FOLDER = "./SummarizationModel"
+SUMMARIZATION_MODEL_WEIGHTS = "pytorch_model.bin"
+SUMMARIZATION_CONFIG = "config.json"
+SUMMARIZATION_VOCAB = "vocab.json"
+SUMMARIZATION_MODEL_WEIGHTS_URL = "https://huggingface.co/facebook/bart-large-cnn/resolve/main/pytorch_model.bin"
+SUMMARIZATION_CONFIG_URL = "https://huggingface.co/facebook/bart-large-cnn/resolve/main/config.json"
+SUMMARIZATION_VOCAB_URL = "https://huggingface.co/facebook/bart-large-cnn/resolve/main/vocab.json"
+SUMMARIZATION_FILES_URLS = [
+    (SUMMARIZATION_MODEL_WEIGHTS_URL, SUMMARIZATION_MODEL_WEIGHTS),
+    (SUMMARIZATION_CONFIG_URL, SUMMARIZATION_CONFIG),
+    (SUMMARIZATION_VOCAB_URL, SUMMARIZATION_VOCAB)
+]
+# TTS
+TTS_FOLDER = "./TTSModel"
+TTS_MODEL_NAME = "vits"
+TTS_MODEL_CONFIG = "config.json"
+TTS_MODEL_WEIGHTS = "pytorch_model.bin"
+TTS_VOCAB = "vocab.json"
+TTS_CONFIG_URL = "https://huggingface.co/kakao-enterprise/vits-vctk/resolve/main/config.json"
+TTS_MODEL_WEIGHTS_URL = "https://huggingface.co/kakao-enterprise/vits-vctk/resolve/main/pytorch_model.bin"
+TTS_VOCAB_URL = "https://huggingface.co/kakao-enterprise/vits-vctk/resolve/main/vocab.json"
+TTS_FILES_URLS = [
+    (TTS_CONFIG_URL, TTS_MODEL_CONFIG),
+    (TTS_MODEL_WEIGHTS_URL, TTS_MODEL_WEIGHTS),
+    (TTS_VOCAB_URL, TTS_VOCAB)
+]
+# STT
+STT_FOLDER = "./STTModel"
+STT_MODEL_NAME = "wav2vec2"
+STT_MODEL_WEIGHTS = "pytorch_model.bin"
+STT_CONFIG = "config.json"
+STT_VOCAB = "vocab.json"
+STT_MODEL_WEIGHTS_URL = "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/pytorch_model.bin"
+STT_CONFIG_URL = "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/config.json"
+STT_VOCAB_URL = "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/vocab.json"
+STT_FILES_URLS = [
+    (STT_MODEL_WEIGHTS_URL, STT_MODEL_WEIGHTS),
+    (STT_CONFIG_URL, STT_CONFIG),
+    (STT_VOCAB_URL, STT_VOCAB)
+]
+# Sentiment Analysis
+SENTIMENT_FOLDER = "./SentimentModel"
+SENTIMENT_MODEL_WEIGHTS = "pytorch_model.bin"
+SENTIMENT_VOCAB = "vocab.json"
+SENTIMENT_CONFIG_FILE = "config.json"
+SENTIMENT_MODEL_WEIGHTS_URL = "https://huggingface.co/climatebert/distilroberta-base-climate-sentiment/resolve/main/pytorch_model.bin"
+SENTIMENT_VOCAB_URL = "https://huggingface.co/climatebert/distilroberta-base-climate-sentiment/resolve/main/vocab.json"
+SENTIMENT_CONFIG_URL = "https://huggingface.co/climatebert/distilroberta-base-climate-sentiment/resolve/main/config.json"
+SENTIMENT_FILES_URLS = [
+    (SENTIMENT_MODEL_WEIGHTS_URL, SENTIMENT_MODEL_WEIGHTS),
+    (SENTIMENT_VOCAB_URL, SENTIMENT_VOCAB),
+    (SENTIMENT_CONFIG_URL, SENTIMENT_CONFIG_FILE)
+]
+# Image Generation (VAE)
+IMAGEGEN_FOLDER = "./ImageGenModel"
+IMAGEGEN_MODEL_WEIGHTS = "diffusion_pytorch_model.bin"
+IMAGEGEN_CONFIG = "config.json"
+IMAGEGEN_MODEL_WEIGHTS_URL = "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/diffusion_pytorch_model.bin"
+IMAGEGEN_CONFIG_URL = "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/config.json"
+IMAGEGEN_FILES_URLS = [
+    (IMAGEGEN_MODEL_WEIGHTS_URL, IMAGEGEN_MODEL_WEIGHTS),
+    (IMAGEGEN_CONFIG_URL, IMAGEGEN_CONFIG)
+]
+# Image to 3D
+IMAGE_TO_3D_FOLDER = "./ImageTo3DModel"
+IMAGE_TO_3D_MODEL_WEIGHTS = "pytorch_model.bin"
+IMAGE_TO_3D_CONFIG = "config.json"
+IMAGE_TO_3D_MODEL_WEIGHTS_URL = "https://huggingface.co/zxhezexin/openlrm-obj-base-1.1/resolve/main/pytorch_model.bin"
+IMAGE_TO_3D_CONFIG_URL = "https://huggingface.co/zxhezexin/openlrm-obj-base-1.1/resolve/main/config.json"
+IMAGE_TO_3D_FILES_URLS = [
+    (IMAGE_TO_3D_MODEL_WEIGHTS_URL, IMAGE_TO_3D_MODEL_WEIGHTS),
+    (IMAGE_TO_3D_CONFIG_URL, IMAGE_TO_3D_CONFIG)
+]
+# Text to Video
+TEXT_TO_VIDEO_FOLDER = "./TextToVideoModel"
+TEXT_TO_VIDEO_MODEL_WEIGHTS = "diffusion_pytorch_model.bin"  # Usado para ambos (Unet y VAE)
+TEXT_TO_VIDEOX_MODEL_WEIGHTS = "diffusion_pytorch_model.fp16.bin"  # Usado para ambos (Unet y VAE)
+TEXT_TO_VIDEO_CONFIG = "config.json"                          # Usado para ambos (Unet y VAE)
+TEXT_TO_VIDEO_VOCAB = "vocab.json"
+TEXT_TO_VIDEO_MODEL_WEIGHTS_URL_UNET = "https://huggingface.co/ali-vilab/text-to-video-ms-1.7b/resolve/main/unet/diffusion_pytorch_model.fp16.bin"
+TEXT_TO_VIDEO_CONFIG_URL_UNET = "https://huggingface.co/ali-vilab/text-to-video-ms-1.7b/resolve/main/unet/config.json"
+TEXT_TO_VIDEO_MODEL_WEIGHTS_URL_VAE = "https://huggingface.co/ali-vilab/text-to-video-ms-1.7b/resolve/main/vae/diffusion_pytorch_model.fp16.bin"
+TEXT_TO_VIDEO_CONFIG_URL_VAE = "https://huggingface.co/ali-vilab/text-to-video-ms-1.7b/resolve/main/vae/config.json"
+TEXT_TO_VIDEO_VOCAB_URL = "https://huggingface.co/ali-vilab/text-to-video-ms-1.7b/resolve/main/tokenizer/vocab.json"
+TEXT_TO_VIDEO_FILES_URLS = [
+    (TEXT_TO_VIDEO_MODEL_WEIGHTS_URL_UNET, TEXT_TO_VIDEO_MODEL_WEIGHTS),
+    (TEXT_TO_VIDEO_MODEL_WEIGHTS_URL_UNET, TEXT_TO_VIDEOX_MODEL_WEIGHTS),
+    (TEXT_TO_VIDEO_CONFIG_URL_UNET, TEXT_TO_VIDEO_CONFIG),
+    (TEXT_TO_VIDEO_MODEL_WEIGHTS_URL_VAE, TEXT_TO_VIDEO_MODEL_WEIGHTS),
+    (TEXT_TO_VIDEO_MODEL_WEIGHTS_URL_VAE, TEXT_TO_VIDEOX_MODEL_WEIGHTS),
+    (TEXT_TO_VIDEO_CONFIG_URL_VAE, TEXT_TO_VIDEO_CONFIG),
+    (TEXT_TO_VIDEO_VOCAB_URL, TEXT_TO_VIDEO_VOCAB),
+]
+# SadTalker
+# ============================================================================
+# Modelos de Restauración para SadTalker (Face Restoration / Super-Resolution)
+# ============================================================================
+# GFPGAN
+GFPGAN_FOLDER = "./GFPGAN"
+GFPGAN_MODEL_FILE = "GFPGANv1.4.pth"
+GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
+# RestoreFormer
+RESTOREFORMER_FOLDER = "./RestoreFormer"
+RESTOREFORMER_MODEL_FILE = "RestoreFormer.pth"
+RESTOREFORMER_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth"
+# CodeFormer
+CODEFORMER_FOLDER = "./CodeFormer"
+CODEFORMER_MODEL_FILE = "codeformer.pth"
+CODEFORMER_URL = "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth"
+# RealESRGAN
+REALESRGAN_FOLDER = "./RealESRGAN"
+REALESRGAN_MODEL_FILE = "RealESRGAN_x2plus.pth"
+REALESRGAN_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
+kp = "https://huggingface.co/usyd-community/vitpose-base-simple/resolve/main/model.safetensors"
+kp_file = "kp_detector.safetensors"
+aud = "https://huggingface.co/vinthony/SadTalker/resolve/main/auido2pose_00140-model.pth"
+aud_file = "auido2pose_00140-model.pth"
+wav = "https://huggingface.co/facebook/wav2vec2-base/resolve/main/pytorch_model.bin"
+wav_file = "wav2vec2.bin"
+gen = "https://huggingface.co/vinthony/SadTalker/resolve/main/wav2lip.pth"
+gen_file = "generator.bin"
+mapx = "https://huggingface.co/vinthony/SadTalker/resolve/main/mapping_00229-model.pth.tar"
+mapx_file = "mapping.pth"
+den = "https://huggingface.co/KwaiVGI/LivePortrait/resolve/main/liveportrait/base_models/motion_extractor.pth"
+den_file = "dense_motion.pth"
+# --- Define constants for new SadTalker models ---
+SADTALKER_KP_FOLDER = "checkpoints"
+SADTALKER_KP_MODEL_FILE = kp_file
+SADTALKER_KP_URL = kp
+SADTALKER_AUD_FOLDER = "checkpoints" # Assuming these go in the main checkpoints folder for SadTalker
+SADTALKER_AUD_MODEL_FILE = aud_file
+SADTALKER_AUD_URL = aud
+SADTALKER_WAV_FOLDER = "checkpoints" # Assuming these go in the main checkpoints folder for SadTalker
+SADTALKER_WAV_MODEL_FILE = wav_file
+SADTALKER_WAV_URL = wav
+SADTALKER_GEN_FOLDER = "checkpoints" # Assuming these go in the main checkpoints folder for SadTalker
+SADTALKER_GEN_MODEL_FILE = gen_file
+SADTALKER_GEN_URL = gen
+SADTALKER_MAPX_FOLDER = "checkpoints" # Assuming these go in the main checkpoints folder for SadTalker
+SADTALKER_MAPX_MODEL_FILE = mapx_file
+SADTALKER_MAPX_URL = mapx
+SADTALKER_DEN_FOLDER = "checkpoints" # Assuming these go in the main checkpoints folder for SadTalker
+SADTALKER_DEN_MODEL_FILE = den_file
+SADTALKER_DEN_URL = den
+# =============================================================================
+# SadTalker
+# =============================================================================
+SADTALKER_CHECKPOINTS_FOLDER = "./checkpoints"
+SADTALKER_CONFIG_FOLDER = "./src/config"

extensions.py ADDED Viewed

	@@ -0,0 +1,252 @@

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import yaml
+from PIL import Image
+from skimage import img_as_ubyte, transform
+import safetensors
+import librosa
+from pydub import AudioSegment
+import imageio
+from scipy.io import loadmat, savemat, wavfile
+import glob
+import tempfile
+from tqdm import tqdm
+import numpy as np
+import math
+import torchvision
+import os
+import re
+import shutil
+from yacs.config import CfgNode as CN
+import requests
+import subprocess
+import cv2
+from collections import OrderedDict
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    if isinstance(imgs, np.ndarray):
+        if imgs.ndim == 3:
+            imgs = imgs[..., np.newaxis]
+        imgs = torch.from_numpy(imgs.transpose((2, 0, 1)))
+    elif isinstance(imgs, Image.Image):
+        imgs = torch.from_numpy(np.array(imgs)).permute(2, 0, 1)
+    else:
+        raise TypeError(f'Type `{type(imgs)}` is not suitable for img2tensor')
+    if bgr2rgb:
+        if imgs.shape[0] == 3:
+            imgs = imgs[[2, 1, 0], :, :]
+    if float32:
+        imgs = imgs.float() / 255.
+    return imgs
+def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
+    if not isinstance(tensor, torch.Tensor):
+        raise TypeError(f'Input tensor should be torch.Tensor, but got {type(tensor)}')
+    tensor = tensor.float().cpu()
+    tensor = tensor.clamp_(*min_max)
+    tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0])
+    output_img = tensor.mul(255).round()
+    output_img = np.transpose(output_img.numpy(), (1, 2, 0))
+    output_img = np.clip(output_img, 0, 255).astype(np.uint8)
+    if rgb2bgr:
+        output_img = cv2.cvtColor(output_img, cv2.COLOR_RGB2BGR)
+    return output_img if out_type == np.uint8 else output_img.astype(out_type) / 255.
+class RealESRGANer():
+    def __init__(self, scale, model_path, model=None, tile=0, tile_pad=10, pre_pad=0, half=False, device=None, gpu_id=None):
+        self.scale = scale
+        self.tile = tile
+        self.tile_pad = tile_pad
+        self.pre_pad = pre_pad
+        self.mod_scale = None
+        self.half = half
+        if device is None:
+            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        else:
+            self.device = device
+        if model is None:
+            model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=scale)
+        if half:
+            model.half()
+        loadnet = torch.load(model_path, map_location=lambda storage, loc: storage)
+        if 'params' in loadnet:
+            model.load_state_dict(loadnet['params'], strict=True)
+        elif 'params_ema' in loadnet:
+            model.load_state_dict(loadnet['params_ema'], strict=True)
+        else:
+            model.load_state_dict(loadnet, strict=True)
+        model.eval()
+        self.model = model.to(self.device)
+    def enhance(self, img, outscale=None, tile=None, tile_pad=None, pre_pad=None, half=None):
+        h_input, w_input = img.shape[0:2]
+        if outscale is None:
+            outscale = self.scale
+        if tile is None:
+            tile = self.tile
+        if tile_pad is None:
+            tile_pad = self.tile_pad
+        if pre_pad is None:
+            pre_pad = self.pre_pad
+        if half is None:
+            half = self.half
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        img_tensor = img2tensor(img)
+        img_tensor = img_tensor.unsqueeze(0).to(self.device)
+        if half:
+            img_tensor = img_tensor.half()
+        mod_scale = self.mod_scale
+        h_pad, w_pad = 0, 0
+        if mod_scale is not None:
+            h_pad, w_pad = int(np.ceil(h_input / mod_scale) * mod_scale - h_input), int(np.ceil(w_input / mod_scale) * mod_scale - w_input)
+            img_tensor = F.pad(img_tensor, (0, w_pad, 0, h_pad), 'reflect')
+        window_size = 256
+        scale = self.scale
+        overlap_ratio = 0.5
+        if w_input * h_input < window_size**2:
+            tile = None
+        if tile is not None and tile > 0:
+            tile_overlap = tile * overlap_ratio
+            sf = scale
+            stride_w = math.ceil(tile - tile_overlap)
+            stride_h = math.ceil(tile - tile_overlap)
+            numW = math.ceil((w_input + tile_overlap) / stride_w)
+            numH = math.ceil((h_input + tile_overlap) / stride_h)
+            paddingW = (numW - 1) * stride_w + tile - w_input
+            paddingH = (numH - 1) * stride_h + tile - h_input
+            padding_bottom = int(max(paddingH, 0))
+            padding_right = int(max(paddingW, 0))
+            padding_left, padding_top = 0, 0
+            img_tensor = F.pad(img_tensor, (padding_left, padding_right, padding_top, padding_bottom), mode='reflect')
+            output_h, output_w = padding_top + h_input * scale + padding_bottom, padding_left + w_input * scale + padding_right
+            output_tensor = torch.zeros([1, 3, output_h, output_w], dtype=img_tensor.dtype, device=self.device)
+            windows = []
+            for row in range(numH):
+                for col in range(numW):
+                    start_x = col * stride_w
+                    start_y = row * stride_h
+                    end_x = min(start_x + tile, img_tensor.shape[3])
+                    end_y = min(start_y + tile, img_tensor.shape[2])
+                    windows.append(img_tensor[:, :, start_y:end_y, start_x:end_x])
+            results = []
+            batch_size = 8
+            for i in range(0, len(windows), batch_size):
+                batch_windows = torch.stack(windows[i:min(i + batch_size, len(windows))], dim=0)
+                with torch.no_grad():
+                    results.append(self.model(batch_windows))
+            results = torch.cat(results, dim=0)
+            count = 0
+            for row in range(numH):
+                for col in range(numW):
+                    start_x = col * stride_w
+                    start_y = row * stride_h
+                    end_x = min(start_x + tile, img_tensor.shape[3])
+                    end_y = min(start_y + tile, img_tensor.shape[2])
+                    out_start_x, out_start_y = start_x * sf, start_y * sf
+                    out_end_x, out_end_y = end_x * sf, end_y * sf
+                    output_tensor[:, :, out_start_y:out_end_y, out_start_x:out_end_x] += results[count][:, :, :end_y * sf - out_start_y, :end_x * sf - out_start_x]
+                    count += 1
+            forward_img = output_tensor[:, :, :h_input * sf, :w_input * sf]
+        else:
+            with torch.no_grad():
+                forward_img = self.model(img_tensor)
+        if half:
+            forward_img = forward_img.float()
+        output_img = tensor2img(forward_img.squeeze(0).clamp_(0, 1))
+        if mod_scale is not None:
+            output_img = output_img[:h_input * self.scale, :w_input * self.scale, ...]
+        output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB)
+        return [output_img, None]
+def save_video_with_watermark(video_frames, audio_path, output_path, watermark_path='./assets/sadtalker_logo.png'):
+    try:
+        watermark = imageio.imread(watermark_path)
+    except FileNotFoundError:
+        watermark = None
+    writer = imageio.get_writer(output_path, fps=25)
+    try:
+        for frame in tqdm(video_frames, 'Generating video'):
+            if watermark is not None:
+                frame_h, frame_w = frame.shape[:2]
+                watermark_h, watermark_w = watermark.shape[:2]
+                if watermark_h > frame_h or watermark_w > frame_w:
+                    watermark = transform.resize(watermark, (frame_h // 4, frame_w // 4))
+                    watermark_h, watermark_w = watermark.shape[:2]
+                start_h = frame_h - watermark_h - 10
+                start_w = frame_w - watermark_w - 10
+                frame[start_h:start_h+watermark_h, start_w:start_w+watermark_w, :] = watermark
+            writer.append_data(img_as_ubyte(frame))
+    except Exception as e:
+        print(f"Error in video writing: {e}")
+    finally:
+        writer.close()
+    if audio_path is not None:
+        try:
+            command = "ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}".format(audio_path, output_path, output_path.replace('.mp4', '_with_audio.mp4'))
+            subprocess.call(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            os.remove(output_path)
+            os.rename(output_path.replace('.mp4', '_with_audio.mp4'), output_path)
+        except Exception as e:
+            print(f"Error adding audio to video: {e}")
+def paste_pic(video_path, pic_path, crop_info, audio_path, output_path):
+    try:
+        y_start, y_end, x_start, x_end, old_size, cropped_size = crop_info[0][0], crop_info[0][1], crop_info[1][0], crop_info[1][1], crop_info[2], crop_info[3]
+        source_image_h, source_image_w = old_size
+        cropped_h, cropped_w = cropped_size
+        delta_h, delta_w = source_image_h - cropped_h, source_image_w - cropped_w
+        box = [x_start, y_start, source_image_w - x_end, source_image_h - y_end]
+        command = "ffmpeg -y -i {} -i {} -filter_complex \"[1]crop=w={}:h={}:x={}:y={},[s];[0][s]overlay=x={}:y={}\" -codec:a copy {}".format(video_path, pic_path, cropped_w, cropped_h, box[0], box[1], box[0], box[1], output_path)
+        subprocess.call(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    except Exception as e:
+        print(f"Error pasting picture to video: {e}")
+def color_transfer_batch(source, target, mode='numpy'):
+    source_np = tensor2img(source)
+    target_np = tensor2img(target)
+    source_lab = cv2.cvtColor(source_np, cv2.COLOR_RGB2LAB).astype(np.float32)
+    target_lab = cv2.cvtColor(target_np, cv2.COLOR_RGB2LAB).astype(np.float32)
+    source_mu = np.mean(source_lab, axis=(0, 1), keepdims=True)
+    source_std = np.std(source_lab, axis=(0, 1), keepdims=True)
+    target_mu = np.mean(target_lab, axis=(0, 1), keepdims=True)
+    target_std = np.std(target_lab, axis=(0, 1), keepdims=True)
+    transfer_lab = (target_lab - target_mu) * (source_std / target_std) + source_mu
+    transfer_rgb = cv2.cvtColor(np.clip(transfer_lab, 0, 255).astype(np.uint8), cv2.COLOR_LAB2RGB)
+    transfer_rgb_tensor = img2tensor(transfer_rgb)
+    return transfer_rgb_tensor.unsqueeze(0).to(source.device)
+def load_video_to_cv2(path, resize=None):
+    video = []
+    try:
+        cap = cv2.VideoCapture(path)
+        if not cap.isOpened():
+            raise Exception("Error opening video stream or file")
+        while(cap.isOpened()):
+            ret, frame = cap.read()
+            if ret:
+                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                if resize is not None:
+                    frame_rgb = cv2.resize(frame_rgb, resize)
+                video.append(frame_rgb)
+            else:
+                break
+        cap.release()
+    except Exception as e:
+        print(f"Error loading video: {e}")
+    return video
+def get_prior_from_bfm(bfm_path):
+    mat_path = os.path.join(bfm_path, 'BFM_prior.mat')
+    C = loadmat(mat_path)
+    pc_tex = torch.tensor(C['pc_tex'].astype(np.float32)).unsqueeze(0)
+    pc_exp = torch.tensor(C['pc_exp'].astype(np.float32)).unsqueeze(0)
+    u_tex = torch.tensor(C['u_tex'].astype(np.float32)).unsqueeze(0)
+    u_exp = torch.tensor(C['u_exp'].astype(np.float32)).unsqueeze(0)
+    prior_coeff = {
+        'pc_tex': pc_tex,
+        'pc_exp': pc_exp,
+        'u_tex': u_tex,
+        'u_exp': u_exp
+    }
+    return prior_coeff

image_to_3d_api.py ADDED Viewed

	@@ -0,0 +1,32 @@

+import os
+import uuid
+from flask import jsonify, send_file, request
+from main import *
+#from main import import image_to_3d_model, device
+from PIL import Image
+import torch
+import numpy as np
+def image_to_3d_func(image_path, output_path="output_3d.obj"):
+    if image_to_3d_model is None:
+        return "Image-to-3D model not initialized."
+    pil_image = Image.open(image_path).convert("RGB")
+    image = torch.tensor(np.array(pil_image)).float().permute(2,0,1).unsqueeze(0) / 255.0
+    image = image.to(device)
+    with torch.no_grad():
+        mesh_obj = image_to_3d_model(image)
+    with open(output_path, 'w') as f:
+        f.write(mesh_obj)
+    return output_path
+def image_to_3d_api():
+    if 'image' not in request.files:
+        return jsonify({"error": "Image file is required"}), 400
+    image_file = request.files['image']
+    temp_image_path = f"temp_image_{uuid.uuid4()}.png"
+    image_file.save(temp_image_path)
+    output_file = image_to_3d_func(temp_image_path)
+    os.remove(temp_image_path)
+    if output_file == "Image-to-3D model not initialized.":
+        return jsonify({"error": "Image to 3D failed"}), 500
+    return send_file(output_file, mimetype="model/obj", as_attachment=True, download_name="output_3d.obj")

imagegen_api.py ADDED Viewed

	@@ -0,0 +1,33 @@

+import os
+from flask import jsonify, send_file, request
+from io import BytesIO
+from PIL import Image
+from main import *
+#from main import import imagegen_model, device
+import torch
+def generate_image(prompt, output_path="output_image.png"):
+    if imagegen_model is None:
+        return "Image generation model not initialized."
+    generator = torch.Generator(device=device).manual_seed(0)
+    image = imagegen_model(
+        prompt,
+        generator=generator,
+    ).images[0]
+    image.save(output_path)
+    return output_path
+def imagegen_api():
+    data = request.get_json()
+    prompt = data.get('prompt')
+    if not prompt:
+        return jsonify({"error": "Prompt is required"}), 400
+    output_file = generate_image(prompt)
+    if output_file == "Image generation model not initialized.":
+        return jsonify({"error": "Image generation failed"}), 500
+    image_io = BytesIO()
+    pil_image = Image.open(output_file)
+    pil_image.save(image_io, 'PNG')
+    image_io.seek(0)
+    return send_file(image_io, mimetype='image/png', as_attachment=True, download_name="output.png")

main.py ADDED Viewed

	@@ -0,0 +1,118 @@

+import threading
+import queue
+import time
+import os
+import nltk
+import re
+import json
+from flask import Flask
+from flask_cors import CORS
+from api import *
+from extensions import *
+from constants import *
+from configs import *
+from tokenxxx import *
+from models import *
+from model_loader import *
+from utils import *
+from background_tasks import generate_and_queue_text, background_training, background_reasoning_queue
+from text_generation import *
+from sadtalker_utils import *
+import torch
+state_dict = None
+enc = None
+config = None
+model_gpt2 = None
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+news_clf = None
+tfidf_vectorizer = None
+text_queue = queue.Queue()
+categories = None
+background_threads = []
+feedback_queue = queue.Queue()
+reasoning_queue = queue.Queue()
+seen_responses = set()
+dialogue_history = []
+vocabulary = set()
+word_to_index = {}
+index_to_word = []
+translation_model = None
+sp = None
+codegen_model = None
+codegen_tokenizer = None
+codegen_vocabulary = None
+codegen_index_to_word = None
+codegen_word_to_index = None
+summarization_model = None
+summarization_vocabulary = set()
+summarization_word_to_index = {}
+summarization_index_to_word = []
+sadtalker_instance = None
+imagegen_model = None
+image_to_3d_model = None
+text_to_video_model = None
+stream_type = "text"
+sentiment_model = None
+stt_model = None
+tts_model = None
+musicgen_model = None
+def load_models():
+    global model_gpt2, enc, translation_model, codegen_model, codegen_tokenizer, codegen_vocabulary, codegen_index_to_word, codegen_word_to_index, summarization_model, imagegen_model, image_to_3d_model, text_to_video_model, sadtalker_instance, sentiment_model, stt_model, tts_model, musicgen_model, checkpoint_path, gfpgan_model_file, restoreformer_model_file, codeformer_model_file, realesrgan_model_file, kp_file, aud_file, wav_file, gen_file, mapx_file, den_file
+    model_gpt2, enc = initialize_gpt2_model(GPT2_FOLDER, {MODEL_FILE: MODEL_URL, ENCODER_FILE: ENCODER_URL, VOCAB_FILE: VOCAB_URL, CONFIG_FILE: GPT2CONFHG})
+    translation_model = initialize_translation_model(TRANSLATION_FOLDER, TRANSLATION_MODEL_FILES_URLS)
+    codegen_model, codegen_tokenizer, codegen_vocabulary, codegen_index_to_word, codegen_word_to_index = initialize_codegen_model(CODEGEN_FOLDER, CODEGEN_FILES_URLS)
+    summarization_model, _, _, _ = initialize_summarization_model(SUMMARIZATION_FOLDER, SUMMARIZATION_FILES_URLS)
+    imagegen_model = initialize_imagegen_model(IMAGEGEN_FOLDER, IMAGEGEN_FILES_URLS)
+    image_to_3d_model = initialize_image_to_3d_model(IMAGE_TO_3D_FOLDER, IMAGE_TO_3D_FILES_URLS)
+    text_to_video_model = initialize_text_to_video_model(TEXT_TO_VIDEO_FOLDER, TEXT_TO_VIDEO_FILES_URLS)
+    sentiment_model = initialize_sentiment_model(SENTIMENT_FOLDER, SENTIMENT_FILES_URLS)
+    stt_model = initialize_stt_model(STT_FOLDER, STT_FILES_URLS)
+    tts_model = initialize_tts_model(TTS_FOLDER, TTS_FILES_URLS)
+    musicgen_model = initialize_musicgen_model(MUSICGEN_FOLDER, MUSICGEN_FILES_URLS)
+class SimpleClassifier(torch.nn.Module):
+    def __init__(self, vocab_size, num_classes):
+        super(SimpleClassifier, self).__init__()
+        self.embedding = torch.nn.Embedding(vocab_size, 128)
+        self.linear = torch.nn.Linear(128, num_classes)
+    def forward(self, x):
+        embedded = self.embedding(x)
+        pooled = torch.mean(embedded, dim=1)
+        return self.linear(pooled)
+def tokenize_text(text):
+    global vocabulary, word_to_index, index_to_word
+    tokens = text.lower().split()
+    for token in tokens:
+        if token not in vocabulary:
+            vocabulary.add(token)
+            word_to_index[token] = len(index_to_word)
+            index_to_word.append(token)
+    return tokens
+def text_to_vector(text):
+    global vocabulary, word_to_index
+    tokens = tokenize_text(text)
+    vector = torch.zeros(len(vocabulary))
+    for token in tokens:
+        if token in word_to_index:
+            vector[word_to_index[token]] += 1
+    return vector
+if __name__ == "__main__":
+    nltk.download('punkt')
+    load_models()
+    categories = ['Category1', 'Category2', 'Category3', 'Category4', 'Category5']
+    import background_tasks
+    background_tasks.categories = categories
+    background_tasks.text_queue = text_queue
+    background_tasks.reasoning_queue = reasoning_queue
+    background_threads.append(threading.Thread(target=generate_and_queue_text, args=('en',), daemon=True))
+    background_threads.append(threading.Thread(target=generate_and_queue_text, args=('es',), daemon=True))
+    background_threads.append(threading.Thread(target=background_training, daemon=True))
+    background_threads.append(threading.Thread(target=background_reasoning_queue, daemon=True))
+    for thread in background_threads:
+        thread.start()
+    app.run(host='0.0.0.0', port=7860)

model_loader.py ADDED Viewed

	@@ -0,0 +1,674 @@

+import os
+import json
+import urllib.request
+import urllib.parse
+import torch
+import hashlib
+from tqdm import tqdm
+from skimage import img_as_ubyte
+from torch import nn
+import torch.nn.functional as F
+import inspect
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+def filter_kwargs(cls, kwargs):
+    sig = inspect.signature(cls.__init__)
+    accepted = set(sig.parameters.keys()) - {"self"}
+    return {k: v for k, v in kwargs.items() if k in accepted}
+def sanitize_filename(name, url=None):
+    for c in '<>:"/\\|?*':
+        name = name.replace(c, '')
+    if not name and url is not None:
+        name = hashlib.md5(url.encode()).hexdigest()
+    return name
+def download_file(url, filepath):
+    d = os.path.dirname(filepath)
+    if d and not os.path.exists(d):
+        os.makedirs(d, exist_ok=True)
+    if not os.path.exists(filepath):
+        def prog(t):
+            last = [0]
+            def inner(n, bs, ts):
+                if ts > 0:
+                    t.total = ts
+                t.update(n * bs - last[0])
+                last[0] = n * bs
+            return inner
+        with tqdm(unit='B', unit_scale=True, unit_divisor=1024, desc=os.path.basename(filepath)) as t:
+            urllib.request.urlretrieve(url, filepath, reporthook=prog(t))
+def download_files(folder, files_spec):
+    if isinstance(files_spec, dict):
+        for fn, url in files_spec.items():
+            fn = sanitize_filename(fn, url)
+            fp = os.path.join(folder, fn)
+            download_file(url, fp)
+    elif isinstance(files_spec, list):
+        for item in files_spec:
+            if isinstance(item, str):
+                url = item
+                parsed = urllib.parse.urlparse(url)
+                fn = os.path.basename(parsed.path)
+                if not fn:
+                    fn = hashlib.md5(url.encode()).hexdigest()
+                fn = sanitize_filename(fn, url)
+            elif isinstance(item, (list, tuple)) and len(item) == 2:
+                url, fn = item
+                fn = sanitize_filename(fn, url)
+            elif isinstance(item, dict) and "filename" in item and "url" in item:
+                fn = sanitize_filename(item["filename"], item["url"])
+                url = item["url"]
+            else:
+                raise ValueError("Invalid file specification")
+            fp = os.path.join(folder, fn)
+            download_file(url, fp)
+    else:
+        raise ValueError("files_spec must be dict or list")
+def read_json(fp):
+    with open(fp, 'r', encoding='utf-8') as f:
+        return json.load(f)
+def get_codegen_tokenizer(vocab_path, merges_path):
+    with open(vocab_path, 'r', encoding='utf-8') as f:
+        vocab = json.load(f)
+    with open(merges_path, 'r', encoding='utf-8') as f:
+        merges = f.read().splitlines()
+    def tokenizer(text):
+        toks = text.split()
+        return [vocab.get(t, 0) for t in toks]
+    return tokenizer
+def simple_tokenizer(text, vocab, max_length=77):
+    toks = text.split()
+    ids = [vocab.get(t, 1) for t in toks]
+    if len(ids) < max_length:
+        ids = ids + [0]*(max_length - len(ids))
+    else:
+        ids = ids[:max_length]
+    return torch.tensor(ids, dtype=torch.long).unsqueeze(0).to(device)
+def load_state_dict_safe(model, loaded_state_dict):
+    model_state = model.state_dict()
+    new_state = {}
+    for key, value in model_state.items():
+        if key in loaded_state_dict and loaded_state_dict[key].shape == value.shape:
+            new_state[key] = loaded_state_dict[key]
+        else:
+            new_state[key] = value
+    model.load_state_dict(new_state, strict=False)
+class GPT2Config:
+    def __init__(self, vocab_size=50257, **kwargs):
+        self.vocab_size = vocab_size
+        self.__dict__.update(kwargs)
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+class MBartConfig:
+    def __init__(self, vocab_size=50265, **kwargs):
+        self.vocab_size = vocab_size
+        self.__dict__.update(kwargs)
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+class CodeGenConfig:
+    def __init__(self, vocab_size=50257, **kwargs):
+        self.vocab_size = vocab_size
+        self.__dict__.update(kwargs)
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+class BartConfig:
+    def __init__(self, vocab_size=50265, **kwargs):
+        self.vocab_size = vocab_size
+        self.__dict__.update(kwargs)
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+class AutoencoderKLConfig:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+class OpenLRMConfig:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+class UNet2DConditionModelConfig:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+class MusicGenConfig:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+class GPT2LMHeadModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(d_model=768, nhead=12)
+        self.transformer = nn.TransformerEncoder(layer, num_layers=12)
+        self.lm_head = nn.Linear(768, config.vocab_size)
+    def forward(self, x):
+        return self.lm_head(self.transformer(x))
+class MBartForConditionalGeneration(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        layer = nn.TransformerEncoderLayer(d_model=768, nhead=12)
+        self.encoder = nn.TransformerEncoder(layer, num_layers=6)
+        dlayer = nn.TransformerDecoderLayer(d_model=768, nhead=12)
+        self.decoder = nn.TransformerDecoder(dlayer, num_layers=6)
+        self.output_layer = nn.Linear(768, config.vocab_size)
+    def forward(self, src, tgt):
+        return self.output_layer(self.decoder(tgt, self.encoder(src)))
+class CodeGenForCausalLM(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        d_model = getattr(config, "d_model", 1024)
+        n_head = getattr(config, "n_head", 16)
+        num_layers = getattr(config, "num_layers", 12)
+        dlayer = nn.TransformerDecoderLayer(d_model=d_model, nhead=n_head)
+        self.transformer_decoder = nn.TransformerDecoder(dlayer, num_layers=num_layers)
+        self.lm_head = nn.Linear(d_model, config.vocab_size)
+    def forward(self, tgt, memory=None):
+        if memory is None:
+            memory = torch.zeros_like(tgt)
+        return self.lm_head(self.transformer_decoder(tgt, memory))
+class BartForConditionalGeneration(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(d_model=768, nhead=12)
+        self.encoder = nn.TransformerEncoder(layer, num_layers=6)
+        dlayer = nn.TransformerDecoderLayer(d_model=768, nhead=12)
+        self.decoder = nn.TransformerDecoder(dlayer, num_layers=6)
+        self.output_layer = nn.Linear(768, config.vocab_size)
+    def forward(self, src, tgt):
+        return self.output_layer(self.decoder(tgt, self.encoder(src)))
+class ResnetBlock(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(32, in_ch)
+        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
+        self.norm2 = nn.GroupNorm(32, out_ch)
+        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
+        self.conv_shortcut = nn.Conv2d(in_ch, out_ch, 1)
+    def forward(self, x):
+        sc = self.conv_shortcut(x)
+        h = F.silu(self.norm1(x))
+        h = self.conv1(h)
+        h = F.silu(self.norm2(h))
+        h = self.conv2(h)
+        return h + sc
+class Downsample(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.conv = nn.Conv2d(in_ch, out_ch, 3, stride=2, padding=1)
+    def forward(self, x):
+        return self.conv(x)
+class DownBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, num_res):
+        super().__init__()
+        self.resnets = nn.ModuleList([ResnetBlock(in_ch if i == 0 else out_ch, out_ch) for i in range(num_res)])
+        self.downsamplers = nn.ModuleList([Downsample(out_ch, out_ch)])
+    def forward(self, x):
+        for r in self.resnets:
+            x = r(x)
+        for ds in self.downsamplers:
+            x = ds(x)
+        return x
+class Upsample(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.conv = nn.ConvTranspose2d(in_ch, out_ch, 4, stride=2, padding=1)
+    def forward(self, x):
+        return self.conv(x)
+class UpBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, num_res):
+        super().__init__()
+        self.resnets = nn.ModuleList([ResnetBlock(in_ch if i == 0 else out_ch, out_ch) for i in range(num_res)])
+        self.upsampler = Upsample(out_ch, out_ch)
+    def forward(self, x):
+        for r in self.resnets:
+            x = r(x)
+        return self.upsampler(x)
+class AttentionBlock(nn.Module):
+    def __init__(self, ch):
+        super().__init__()
+        self.norm = nn.GroupNorm(32, ch)
+        self.query = nn.Conv2d(ch, ch, 1)
+        self.key = nn.Conv2d(ch, ch, 1)
+        self.value = nn.Conv2d(ch, ch, 1)
+        self.proj_attn = nn.Conv2d(ch, ch, 1)
+    def forward(self, x):
+        b, c, h, w = x.shape
+        xn = self.norm(x)
+        q = self.query(xn).view(b, c, -1).permute(0, 2, 1)
+        k = self.key(xn).view(b, c, -1)
+        v = self.value(xn).view(b, c, -1).permute(0, 2, 1)
+        attn = torch.softmax(torch.bmm(q, k) / (c ** 0.5), dim=-1)
+        out = torch.bmm(attn, v).permute(0, 2, 1).view(b, c, h, w)
+        return x + self.proj_attn(out)
+class Encoder(nn.Module):
+    def __init__(self, in_ch=3, base_ch=128, latent_ch=4):
+        super().__init__()
+        self.conv_in = nn.Conv2d(in_ch, base_ch, 3, padding=1)
+        self.down_blocks = nn.ModuleList([
+            DownBlock(base_ch, base_ch, 2),
+            DownBlock(base_ch, base_ch * 2, 2),
+            DownBlock(base_ch * 2, base_ch * 4, 2),
+            DownBlock(base_ch * 4, base_ch * 4, 2)
+        ])
+        self.mid_block = nn.ModuleList([
+            ResnetBlock(base_ch * 4, base_ch * 4),
+            AttentionBlock(base_ch * 4),
+            ResnetBlock(base_ch * 4, base_ch * 4)
+        ])
+        self.conv_norm_out = nn.GroupNorm(32, base_ch * 4)
+        self.conv_out = nn.Conv2d(base_ch * 4, latent_ch * 2, 3, padding=1)
+        self.quant_conv = nn.Conv2d(latent_ch * 2, latent_ch, 1)
+    def forward(self, x):
+        x = self.conv_in(x)
+        for blk in self.down_blocks:
+            x = blk(x)
+        for m in self.mid_block:
+            x = m(x)
+        x = self.conv_norm_out(x)
+        x = self.conv_out(x)
+        return self.quant_conv(x)
+class Decoder(nn.Module):
+    def __init__(self, out_ch=3, base_ch=128, latent_ch=4):
+        super().__init__()
+        self.post_quant_conv = nn.Conv2d(latent_ch, latent_ch * 2, 1)
+        self.conv_in = nn.Conv2d(latent_ch, base_ch * 4, 3, padding=1)
+        self.mid_block = nn.ModuleList([
+            ResnetBlock(base_ch * 4, base_ch * 4),
+            AttentionBlock(base_ch * 4),
+            ResnetBlock(base_ch * 4, base_ch * 4)
+        ])
+        self.up_blocks = nn.ModuleList([
+            UpBlock(base_ch * 4, base_ch * 4, 3),
+            UpBlock(base_ch * 4, base_ch * 2, 3),
+            UpBlock(base_ch * 2, base_ch, 3),
+            UpBlock(base_ch, base_ch, 3)
+        ])
+        self.conv_norm_out = nn.GroupNorm(32, base_ch)
+        self.conv_out = nn.Conv2d(base_ch, out_ch, 3, padding=1)
+    def forward(self, x):
+        x = self.post_quant_conv(x)
+        x = self.conv_in(x)
+        for m in self.mid_block:
+            x = m(x)
+        for up in self.up_blocks:
+            x = up(x)
+        x = self.conv_norm_out(x)
+        return self.conv_out(x)
+class AutoencoderKL(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        in_ch = config.get("in_channels", 3) if isinstance(config, dict) else config.__dict__.get("in_channels", 3)
+        out_ch = config.get("out_channels", 3) if isinstance(config, dict) else config.__dict__.get("out_channels", 3)
+        base_ch = config.get("base_channels", 128) if isinstance(config, dict) else config.__dict__.get("base_channels", 128)
+        latent_ch = config.get("latent_channels", 4) if isinstance(config, dict) else config.__dict__.get("latent_channels", 4)
+        self.encoder = Encoder(in_ch, base_ch, latent_ch)
+        self.decoder = Decoder(out_ch, base_ch, latent_ch)
+    def forward(self, x):
+        return self.decoder(self.encoder(x))
+    def decode(self, x):
+        return self.decoder(x)
+class TransformerBlock(nn.Module):
+    def __init__(self, embed_dim, num_heads):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        hidden_dim = embed_dim * 4
+        self.mlp = nn.Sequential(
+            nn.Linear(embed_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, embed_dim)
+        )
+    def forward(self, x):
+        res = x
+        x = self.norm1(x)
+        x = x.transpose(0, 1)
+        attn, _ = self.attn(x, x, x)
+        x = attn.transpose(0, 1)
+        x = res + x
+        return x + self.mlp(self.norm2(x))
+class VisionTransformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        if isinstance(config, dict):
+            self.img_size = config.get("img_size", 592)
+            self.patch_size = config.get("patch_size", 16)
+            self.embed_dim = config.get("hidden_size", 768)
+            depth = config.get("depth", 12)
+            num_heads = config.get("num_heads", 12)
+        else:
+            self.img_size = config.__dict__.get("img_size", 592)
+            self.patch_size = config.__dict__.get("patch_size", 16)
+            self.embed_dim = config.__dict__.get("hidden_size", 768)
+            depth = config.__dict__.get("depth", 12)
+            num_heads = config.__dict__.get("num_heads", 12)
+        num_patches = (self.img_size // self.patch_size) ** 2
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, self.embed_dim))
+        self.patch_embed = nn.Conv2d(3, self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size)
+        self.blocks = nn.ModuleList([TransformerBlock(self.embed_dim, num_heads) for _ in range(depth)])
+        self.norm = nn.LayerNorm(self.embed_dim)
+        self.register_tokens = nn.Parameter(torch.zeros(1, 4, self.embed_dim))
+        self._init_weights()
+    def _init_weights(self):
+        nn.init.normal_(self.cls_token, std=0.02)
+        nn.init.normal_(self.pos_embed, std=0.02)
+    def forward(self, x):
+        x = self.patch_embed(x)
+        x = x.flatten(2).transpose(1, 2)
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embed
+        for blk in self.blocks:
+            x = blk(x)
+        return self.norm(x)[:, 0]
+class OpenLRM(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.encoder = nn.ModuleDict({"model": VisionTransformer(config)})
+        hidden = config.get("hidden_size", 768) if isinstance(config, dict) else config.__dict__.get("hidden_size", 768)
+        self.linear = nn.Linear(hidden, hidden)
+    def forward(self, x):
+        return self.linear(self.encoder["model"](x))
+class VideoUNet(nn.Module):
+    def __init__(self, in_ch=4, out_ch=4, features=None):
+        super().__init__()
+        if features is None:
+            features = [64, 128, 256]
+        self.encoder = nn.ModuleList()
+        self.pool = nn.MaxPool3d(2, 2)
+        self.decoder = nn.ModuleList()
+        for f in features:
+            self.encoder.append(nn.Sequential(
+                nn.Conv3d(in_ch, f, 3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.Conv3d(f, f, 3, padding=1),
+                nn.ReLU(inplace=True)
+            ))
+            in_ch = f
+        for f in reversed(features):
+            self.decoder.append(nn.Sequential(
+                nn.Conv3d(f * 2, f, 3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.Conv3d(f, f, 3, padding=1),
+                nn.ReLU(inplace=True)
+            ))
+        self.final_conv = nn.Conv3d(features[0], out_ch, 1)
+    def forward(self, x, t, encoder_hidden_states):
+        skips = []
+        for enc in self.encoder:
+            x = enc(x)
+            skips.append(x)
+            x = self.pool(x)
+        for dec in self.decoder:
+            skip = skips.pop()
+            x = F.interpolate(x, scale_factor=2, mode='trilinear', align_corners=False)
+            x = torch.cat([x, skip], dim=1)
+            x = dec(x)
+        return self.final_conv(x)
+class SentimentClassifierModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.classifier = nn.Sequential(
+            nn.Linear(768, 256),
+            nn.ReLU(),
+            nn.Linear(256, 2)
+        )
+    def forward(self, x):
+        return self.classifier(x)
+class STTModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(768, 512),
+            nn.ReLU(),
+            nn.Linear(512, 768)
+        )
+    def forward(self, x):
+        return self.net(x)
+class TTSModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(768, 512),
+            nn.ReLU(),
+            nn.Linear(512, 768)
+        )
+    def forward(self, x):
+        return self.net(x)
+class MusicGenModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(d_model=768, nhead=12)
+        self.transformer = nn.TransformerEncoder(layer, num_layers=12)
+        self.linear = nn.Linear(768, 768)
+    def forward(self, x):
+        return self.linear(self.transformer(x))
+class SimpleTextEncoder(nn.Module):
+    def __init__(self, vocab_size=10000, embed_dim=768, max_length=77):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.max_length = max_length
+    def forward(self, text_tokens):
+        return self.embedding(text_tokens)
+class DiffusionScheduler:
+    def __init__(self, steps):
+        self.steps = steps
+        self.betas = torch.linspace(0.1, 0.001, steps=steps).to(device)
+    def step(self, noise, t, sample):
+        beta = self.betas[t]
+        return sample - beta * noise
+class VideoOutput:
+    def __init__(self, frames):
+        self.frames = [img_as_ubyte(frame) for frame in frames[0]]
+class VideoPipeline(nn.Module):
+    def __init__(self, unet, vae, text_encoder, vocab):
+        super().__init__()
+        self.unet = unet
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.vocab = vocab
+    def forward(self, prompt: str, steps: int = 25, num_frames: int = 24):
+        token_ids = simple_tokenizer(prompt, self.vocab)
+        text_emb = self.text_encoder(token_ids)
+        latent = torch.randn((1, 4, num_frames, 64, 64), device=device).half()
+        sched = DiffusionScheduler(steps)
+        for t in range(steps):
+            noise = self.unet(latent, t, text_emb)
+            latent = sched.step(noise, t, latent)
+        frames = self.vae.decode(latent / 0.18215)
+        frames = frames.clamp(0, 1).float().cpu().permute(0, 2, 3, 4, 1).numpy()
+        return VideoOutput(frames)
+def initialize_gpt2_model(folder, files):
+    download_files(folder, files)
+    config = GPT2Config()
+    model = GPT2LMHeadModel(config).to(device)
+    sd = torch.load(os.path.join(folder, sanitize_filename("gpt2-pytorch_model.bin")), map_location=device)
+    load_state_dict_safe(model, sd)
+    model.eval()
+    enc = read_json(os.path.join(folder, sanitize_filename("encoder.json")))
+    return model, enc
+def initialize_translation_model(folder, files):
+    download_files(folder, files)
+    config = MBartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model = MBartForConditionalGeneration(config).to(device)
+    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(model, sd)
+    model.eval()
+    vp = os.path.join(folder, "vocab.json")
+    if os.path.exists(vp):
+        vocab = read_json(vp)
+        model.tokenizer = lambda txt: [vocab.get(t, 0) for t in txt.split()]
+    else:
+        model.tokenizer = lambda txt: txt
+    model.config.lang_code_to_id = {'en_XX': 0, 'es_XX': 1}
+    return model
+def initialize_codegen_model(folder, files):
+    download_files(folder, files)
+    config = CodeGenConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model = CodeGenForCausalLM(config).to(device)
+    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(model, sd)
+    model.eval()
+    tok = get_codegen_tokenizer(os.path.join(folder, "vocab.json"), os.path.join(folder, "merges.txt"))
+    vocab = read_json(os.path.join(folder, "vocab.json"))
+    idx2w = {v: k for k, v in vocab.items()}
+    model.tokenizer = tok
+    return model, tok, vocab, idx2w, vocab
+def initialize_summarization_model(folder, files):
+    download_files(folder, files)
+    config = BartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model = BartForConditionalGeneration(config).to(device)
+    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(model, sd)
+    model.eval()
+    vp = os.path.join(folder, "vocab.json")
+    if os.path.exists(vp):
+        vocab_json = read_json(vp)
+        vocab = set(vocab_json.keys())
+        return model, vocab, vocab_json, {v: k for k, v in vocab_json.items()}
+    return model, None, None, None
+def initialize_imagegen_model(folder, files):
+    download_files(folder, files)
+    config = AutoencoderKLConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    vae = AutoencoderKL(config).to(device)
+    sd = torch.load(os.path.join(folder, "diffusion_pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(vae, sd)
+    vae.eval()
+    return vae
+def initialize_image_to_3d_model(folder, files):
+    download_files(folder, files)
+    config = OpenLRMConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model3d = OpenLRM(config).to(device)
+    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(model3d, sd)
+    model3d.eval()
+    return model3d
+def initialize_text_to_video_model(folder, files):
+    download_files(folder, files)
+    unet_cfg = read_json(os.path.join(folder, "config.json"))
+    unet_cfg = filter_kwargs(VideoUNet, unet_cfg)
+    unet = VideoUNet(**unet_cfg).half().to(device)
+    sd_unet = torch.load(os.path.join(folder, "diffusion_pytorch_model.fp16.bin"), map_location=device)
+    load_state_dict_safe(unet, sd_unet)
+    unet.eval()
+    vae_cfg = read_json(os.path.join(folder, "config.json"))
+    vae_cfg = filter_kwargs(AutoencoderKL, vae_cfg)
+    vae = AutoencoderKL(vae_cfg).half().to(device)
+    sd_vae = torch.load(os.path.join(folder, "diffusion_pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(vae, sd_vae)
+    vae.eval()
+    vp = os.path.join(folder, "vocab.json")
+    text_vocab = read_json(vp) if os.path.exists(vp) else {}
+    te_path = os.path.join(folder, "text_encoder.bin")
+    if os.path.exists(te_path):
+        text_encoder = SimpleTextEncoder(vocab_size=(max(text_vocab.values())+1) if text_vocab else 10000, embed_dim=768, max_length=77).to(device)
+        sd_te = torch.load(te_path, map_location=device)
+        load_state_dict_safe(text_encoder, sd_te)
+    else:
+        text_encoder = SimpleTextEncoder(vocab_size=(max(text_vocab.values())+1) if text_vocab else 10000, embed_dim=768, max_length=77).to(device)
+    text_encoder.eval()
+    return VideoPipeline(unet, vae, text_encoder, text_vocab)
+def initialize_sentiment_model(folder, files):
+    download_files(folder, files)
+    config = BartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model = SentimentClassifierModel(config).to(device)
+    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(model, sd)
+    model.eval()
+    vp = os.path.join(folder, "vocab.json")
+    if os.path.exists(vp):
+        read_json(vp)
+    return model
+def initialize_stt_model(folder, files):
+    download_files(folder, files)
+    config = BartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model = STTModel(config).to(device)
+    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(model, sd)
+    model.eval()
+    vp = os.path.join(folder, "vocab.json")
+    if os.path.exists(vp):
+        read_json(vp)
+    return model
+def initialize_tts_model(folder, files):
+    download_files(folder, files)
+    config = BartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model = TTSModel(config).to(device)
+    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(model, sd)
+    model.eval()
+    vp = os.path.join(folder, "vocab.json")
+    if os.path.exists(vp):
+        read_json(vp)
+    return model
+def initialize_musicgen_model(folder, files):
+    download_files(folder, files)
+    config = MusicGenConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model = MusicGenModel(config).to(device)
+    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
+    load_state_dict_safe(model, sd)
+    model.eval()
+    return model

models.py ADDED Viewed

	@@ -0,0 +1,96 @@

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import copy
+#from configs import GPT2Config, MBartConfig, CodeGenConfig, SummarizationConfig, OpenLRMConfig, UNet2DConditionModelConfig, AutoencoderKLConfig, BartConfig, MusicGenConfig
+from configs import *
+#from extensions import gelu, LayerNorm, Conv1D, Attention, MLP, Block, GPT2Model, GPT2LMHead, MBartEncoderLayer, MBartDecoderLayer, MBartEncoder, MBartDecoder, MBartModel, MBartForConditionalGeneration, CodeGenAttention, CodeGenBlock, CodeGenModel, CodeGenForCausalLM, SummarizationModel, OpenLRM, OpenLRMLayer, OpenLRMAttention, OpenLRMFeedForward, AutoencoderKL, Encoder_, Decoder_, DownBlock, UpBlock, ResnetBlock, MidBlock, Downsample2D, Upsample2D, UNet2DConditionModel, UNetMidBlock2DConditionModel, UNetDownBlock2DConditionModel, UNetUpBlock2DConditionModel, ResnetBlock2D, CrossAttentionBlock2D, CrossAttention, SimpleClassifier
+from extensions import *
+class SentimentClassifierModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.lstm = nn.LSTM(config.d_model, config.d_model, batch_first=True, bidirectional=True)
+        self.fc = nn.Linear(config.d_model * 2, 3)
+    def forward(self, input_ids):
+        embedded = self.embedding(input_ids)
+        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, lengths=[input_ids.size(1)]*input_ids.size(0), batch_first=True, enforce_sorted=False)
+        packed_output, _ = self.lstm(packed_embedded)
+        output, _ = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)
+        pooled = output[:, -1, :]
+        logits = self.fc(pooled)
+        return logits
+class STTModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.conv1 = nn.Conv1d(1, 16, kernel_size=3, stride=2, padding=1)
+        self.relu1 = nn.ReLU()
+        self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.conv2 = nn.Conv1d(16, 32, kernel_size=3, padding=1)
+        self.relu2 = nn.ReLU()
+        self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.lstm = nn.LSTM(32 * (config.max_position_embeddings // 8), 128, batch_first=True, bidirectional=True)
+        self.fc = nn.Linear(128 * 2, config.vocab_size)
+    def forward(self, audio_data):
+        x = self.pool1(self.relu1(self.conv1(audio_data.unsqueeze(1))))
+        x = self.pool2(self.relu2(self.conv2(x)))
+        x = x.transpose(1, 2).contiguous()
+        x = x.view(x.size(0), -1, x.size(2))
+        packed_output = nn.utils.rnn.pack_padded_sequence(x, lengths=[x.size(1)]*x.size(0), batch_first=True, enforce_sorted=False)
+        packed_output, _ = self.lstm(packed_output)
+        output, _ = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)
+        logits = self.fc(output)
+        return logits
+class TTSModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.lstm = nn.LSTM(config.d_model, config.d_model, batch_first=True, bidirectional=True)
+        self.fc = nn.Linear(config.d_model * 2, 1)
+        self.sigmoid = nn.Sigmoid()
+    def forward(self, input_ids):
+        embedded = self.embedding(input_ids)
+        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, lengths=[input_ids.size(1)]*input_ids.size(0), batch_first=True, enforce_sorted=False)
+        packed_output, _ = self.lstm(packed_embedded)
+        output, _ = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)
+        logits = self.fc(output)
+        audio = self.sigmoid(logits)
+        return audio
+class MusicGenModel(nn.Module):
+    def __init__(self, config: MusicGenConfig):
+        super().__init__()
+        self.config = config
+        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.transformer_layers = nn.ModuleList([CodeGenBlock(config) for _ in range(config.num_hidden_layers)])
+        self.fc_out = nn.Linear(config.hidden_size, config.vocab_size)
+    def forward(self, input_ids):
+        embedded_tokens = self.embedding(input_ids)
+        hidden_states = embedded_tokens
+        for layer in self.transformer_layers:
+            hidden_states = layer(hidden_states)
+        logits = self.fc_out(hidden_states)
+        return logits
+    def sample(self, attributes, sample_rate, duration):
+        input_tokens = torch.randint(0, self.config.vocab_size, (1, 1), dtype=torch.long).to(device)
+        audio_output = []
+        num_steps = int(duration * sample_rate / 1024)
+        for _ in tqdm(range(num_steps), desc="Generating music"):
+            logits = self.forward(input_tokens)
+            predicted_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
+            audio_output.append(predicted_token.cpu())
+            input_tokens = torch.cat((input_tokens, predicted_token), dim=1)
+        audio_output = torch.cat(audio_output, dim=1).float()
+        return audio_output

musicgen_api.py ADDED Viewed

	@@ -0,0 +1,35 @@

+from flask import jsonify, send_file, request
+from main import *
+#from main import import musicgen_model, device
+import torch
+import soundfile as sf
+import numpy as np
+import io
+def generate_music(prompt, output_path="output_music.wav"):
+    if musicgen_model is None:
+        return "Music generation model not initialized."
+    attributes = [prompt]
+    sample_rate = 32000
+    duration = 8
+    audio_values = musicgen_model.sample(
+        attributes=attributes,
+        sample_rate=sample_rate,
+        duration=duration,
+    )
+    output_audio = audio_values.cpu().numpy().squeeze()
+    sf.write(output_path, output_audio, sample_rate)
+    return output_path
+def musicgen_api():
+    data = request.get_json()
+    prompt = data.get('prompt')
+    if not prompt:
+        return jsonify({"error": "Prompt is required"}), 400
+    output_file = generate_music(prompt)
+    if output_file == "Music generation model not initialized.":
+        return jsonify({"error": "Music generation failed"}), 500
+    with open(output_file, 'rb') as f:
+        audio_content = f.read()
+    return send_file(io.BytesIO(audio_content), mimetype="audio/wav", as_attachment=True, download_name="output.wav")

requirements.txt ADDED Viewed

	@@ -0,0 +1,40 @@

+accelerate
+retry
+asyncio
+basicsr
+beautifulsoup4
+bs4
+opencv-python
+deep-translator
+duckduckgo-search
+fastapi
+faker
+flask
+flask-cors
+facexlib
+ffmpeg-python
+gfpgan
+imageio
+imageio-ffmpeg
+langdetect
+librosa
+nltk
+numpy
+Pillow
+pydub
+pytorch-lightning
+PyYAML
+retry
+safetensors
+scikit-learn
+scipy
+scikit-image
+soundfile
+torch
+torchaudio
+torchvision
+tqdm
+wget
+yacs
+numba
+librosa

sadtalker_api.py ADDED Viewed

	@@ -0,0 +1,202 @@

+import os
+import tempfile
+import uuid
+import asyncio
+import shutil
+import requests
+from urllib.parse import urlparse
+from fastapi import FastAPI, UploadFile, File, HTTPException, Form, WebSocket
+from fastapi.responses import JSONResponse
+#from fastapi.middleware.cors import CORSMiddleware
+from fastapi import APIRouter
+from extensions import *
+from main import *
+#from main import import sadtalker_instance
+from tts_api import *
+from sadtalker_utils import *
+import base64
+from stt_api import  *
+from text_generation import *
+router = APIRouter()
+@router.post("/sadtalker")
+async def create_video(
+    source_image: str = Form(None),
+    source_image_file: UploadFile = File(None),
+    driven_audio: str = Form(None),
+    driven_audio_file: UploadFile = File(None),
+    preprocess: str = Form('crop'),
+    still_mode: bool = Form(False),
+    use_enhancer: bool = Form(False),
+    batch_size: int = Form(1),
+    size: int = Form(256),
+    pose_style: int = Form(0),
+    exp_scale: float = Form(1.0),
+    use_ref_video: bool = Form(False),
+    ref_video: str = Form(None),
+    ref_video_file: UploadFile = File(None),
+    ref_info: str = Form(None),
+    use_idle_mode: bool = Form(False),
+    length_of_audio: int = Form(0),
+    use_blink: bool = Form(True),
+    checkpoint_dir: str = Form('checkpoints'),
+    config_dir: str = Form('src/config'),
+    old_version: bool = Form(False),
+    tts_text: str = Form(None),
+    tts_lang: str = Form('en'),
+):
+    if source_image_file and source_image:
+         raise HTTPException(status_code=400, detail="source_image and source_image_file cannot be both not None")
+    if driven_audio and driven_audio_file:
+        raise HTTPException(status_code=400, detail="driven_audio and driven_audio_file cannot be both not None")
+    if ref_video and ref_video_file:
+        raise HTTPException(status_code=400, detail="ref_video and ref_video_file cannot be both not None")
+    tmp_source_image = None
+    if source_image_file:
+            tmp_source_image = tempfile.NamedTemporaryFile(suffix=os.path.splitext(source_image_file.filename)[1], delete=False)
+            content = await source_image_file.read()
+            tmp_source_image.write(content)
+            source_image_path = tmp_source_image.name
+    elif source_image:
+            if urlparse(source_image).scheme in ["http", "https"]:
+                response = requests.get(source_image, stream=True)
+                response.raise_for_status()
+                with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_source_image:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        tmp_source_image.write(chunk)
+                    source_image_path = tmp_source_image.name
+            else:
+                source_image_path = source_image
+    else:
+        raise HTTPException(status_code=400, detail="source_image not provided")
+    tmp_driven_audio = None
+    if driven_audio_file:
+         tmp_driven_audio = tempfile.NamedTemporaryFile(suffix=os.path.splitext(driven_audio_file.filename)[1], delete=False)
+         content = await driven_audio_file.read()
+         tmp_driven_audio.write(content)
+         driven_audio_path = tmp_driven_audio.name
+    elif driven_audio:
+         if urlparse(driven_audio).scheme in ["http", "https"]:
+             response = requests.get(driven_audio, stream=True)
+             response.raise_for_status()
+             with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_driven_audio:
+                 for chunk in response.iter_content(chunk_size=8192):
+                     tmp_driven_audio.write(chunk)
+                 driven_audio_path = tmp_driven_audio.name
+         else:
+            driven_audio_path = driven_audio
+    else:
+        driven_audio_path = None
+    tmp_ref_video = None
+    if ref_video_file:
+        tmp_ref_video = tempfile.NamedTemporaryFile(suffix=os.path.splitext(ref_video_file.filename)[1], delete=False)
+        content = await ref_video_file.read()
+        tmp_ref_video.write(content)
+        ref_video_path = tmp_ref_video.name
+    elif ref_video:
+          if urlparse(ref_video).scheme in ["http", "https"]:
+              response = requests.get(ref_video, stream=True)
+              response.raise_for_status()
+              with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_ref_video:
+                  for chunk in response.iter_content(chunk_size=8192):
+                      tmp_ref_video.write(chunk)
+                  ref_video_path = tmp_ref_video.name
+          else:
+            ref_video_path = ref_video
+    else:
+        ref_video_path=None
+    try:
+        loop = asyncio.get_running_loop()
+        output_path = await loop.run_in_executor(None, sadtalker_instance.test,
+            source_image_path,
+            driven_audio_path,
+            preprocess,
+            still_mode,
+            use_enhancer,
+            batch_size,
+            size,
+            pose_style,
+            exp_scale,
+            use_ref_video,
+            ref_video_path,
+            ref_info,
+            use_idle_mode,
+            length_of_audio,
+            use_blink,
+            './results/',
+            tts_text=tts_text,
+            tts_lang=tts_lang,
+        )
+        return {"video_url": output_path}
+    except Exception as e:
+         raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        if tmp_source_image:
+           os.remove(tmp_source_image.name)
+        if tmp_driven_audio:
+            os.remove(tmp_driven_audio.name)
+        if tmp_ref_video:
+            os.remove(tmp_ref_video.name)
+@router.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    tts_model = TTSTalker()
+    try:
+        while True:
+             data = await websocket.receive_json()
+             text = data.get("text")
+             audio_base64 = data.get("audio")
+             if text:
+                audio_path = await asyncio.get_running_loop().run_in_executor(None,  tts_model.test, text)
+             elif audio_base64:
+                try:
+                  audio_bytes = base64.b64decode(audio_base64)
+                  tmp_audio_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+                  tmp_audio_file.write(audio_bytes)
+                  audio_path = tmp_audio_file.name
+                  transcription_text_file = speech_to_text_func(tmp_audio_file.name)
+                  with open(transcription_text_file, 'r') as f:
+                      transcription_text = f.read()
+                  response_stream = perform_reasoning_stream(f"respond to this sentence in 10 words or less {transcription_text}", 0.7, 40, 0.0, 1.2)
+                  response_text = ""
+                  for chunk in response_stream:
+                      if chunk == "<END_STREAM>":
+                          break
+                      response_text += chunk
+                  audio_path = await asyncio.get_running_loop().run_in_executor(None,  tts_model.test, response_text)
+                except Exception as e:
+                    await websocket.send_json({"error":str(e)})
+                    continue
+                finally:
+                   if 'tmp_audio_file' in locals() and tmp_audio_file:
+                       os.remove(tmp_audio_file.name)
+             else:
+                  continue
+             source_image_path = './examples/source_image/cyarh.png'
+             ref_video_path='./examples/driven_video/vid_xdd.mp4'
+             loop = asyncio.get_running_loop()
+             output = await loop.run_in_executor(None, sadtalker_instance.test,
+                source_image_path,
+                audio_path,
+                'full',
+                 True,
+                 True,
+                 1,
+                 256,
+                 0,
+                 1,
+                 True,
+                 ref_video_path,
+                 "pose+blink",
+                 False,
+                 0,
+                 True,
+                './results/'
+             )
+             await websocket.send_json({"video_url": output})
+    except Exception as e:
+         print(e)
+         await websocket.send_json({"error":str(e)})

sadtalker_utils.py ADDED Viewed

	@@ -0,0 +1,866 @@

+import os
+import shutil
+import uuid
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import yaml
+from PIL import Image
+from skimage import img_as_ubyte, transform
+import safetensors
+import librosa
+from pydub import AudioSegment
+import imageio
+from scipy import signal
+from scipy.io import loadmat, savemat, wavfile
+import glob
+import tempfile
+from tqdm import tqdm
+import math
+import torchaudio
+import urllib.request
+REALESRGAN_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
+CODEFORMER_URL = "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth"
+RESTOREFORMER_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth"
+GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
+kp_url = "https://huggingface.co/usyd-community/vitpose-base-simple/resolve/main/model.safetensors"
+kp_file = "kp_detector.safetensors"
+aud_url = "https://huggingface.co/vinthony/SadTalker/resolve/main/auido2pose_00140-model.pth"
+aud_file = "auido2pose_00140-model.pth"
+wav_url = "https://huggingface.co/facebook/wav2vec2-base/resolve/main/pytorch_model.bin"
+wav_file = "wav2vec2.pth"
+gen_url = "https://huggingface.co/vinthony/SadTalker/resolve/main/wav2lip.pth"
+gen_file = "generator.pth"
+mapx_url = "https://huggingface.co/vinthony/SadTalker/resolve/main/mapping_00229-model.pth.tar"
+mapx_file = "mapping.pth"
+den_url = "https://huggingface.co/KwaiVGI/LivePortrait/resolve/main/liveportrait/base_models/motion_extractor.pth"
+den_file = "dense_motion.pth"
+def download_model(url, filename, checkpoint_dir):
+    if not os.path.exists(os.path.join(checkpoint_dir, filename)):
+        print(f"Downloading {filename}...")
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        urllib.request.urlretrieve(url, os.path.join(checkpoint_dir, filename))
+        print(f"{filename} downloaded.")
+    else:
+        print(f"{filename} already exists.")
+def mp3_to_wav_util(mp3_filename, wav_filename, frame_rate):
+    AudioSegment.from_file(mp3_filename).set_frame_rate(frame_rate).export(wav_filename, format="wav")
+def load_wav_util(path, sr):
+    return librosa.core.load(path, sr=sr)[0]
+def save_wav_util(wav, path, sr):
+    wav *= 32767 / max(0.01, np.max(np.abs(wav)))
+    wavfile.write(path, sr, wav.astype(np.int16))
+class OcclusionAwareKPDetector(nn.Module):
+    def __init__(self, kp_channels, num_kp, num_dilation_blocks, dropout_rate):
+        super(OcclusionAwareKPDetector, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(64, num_kp, kernel_size=3, padding=1)
+    def forward(self, x):
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.conv2(x)
+        kp = {'value': x.view(x.size(0), -1)}
+        return kp
+class Wav2Vec2Model(nn.Module):
+    def __init__(self):
+        super(Wav2Vec2Model, self).__init__()
+        self.conv = nn.Conv1d(1, 64, kernel_size=10, stride=5, padding=5)
+        self.bn = nn.BatchNorm1d(64)
+        self.relu = nn.ReLU()
+        self.fc = nn.Linear(64, 2048)
+    def forward(self, audio):
+        x = audio.unsqueeze(1)
+        x = self.relu(self.bn(self.conv(x)))
+        x = torch.mean(x, dim=-1)
+        x = self.fc(x)
+        return x
+class AudioCoeffsPredictor(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super(AudioCoeffsPredictor, self).__init__()
+        self.linear = nn.Linear(input_dim, output_dim)
+    def forward(self, audio_embedding):
+        return self.linear(audio_embedding)
+class MappingNet(nn.Module):
+    def __init__(self, num_coeffs, num_layers, hidden_dim):
+        super(MappingNet, self).__init__()
+        layers = []
+        input_dim = num_coeffs * 2
+        for _ in range(num_layers):
+            layers.append(nn.Linear(input_dim, hidden_dim))
+            layers.append(nn.ReLU())
+            input_dim = hidden_dim
+        layers.append(nn.Linear(hidden_dim, num_coeffs))
+        self.net = nn.Sequential(*layers)
+    def forward(self, x):
+        return self.net(x)
+class DenseMotionNetwork(nn.Module):
+    def __init__(self, num_kp, num_channels, block_expansion, num_blocks, max_features):
+        super(DenseMotionNetwork, self).__init__()
+        self.conv1 = nn.Conv2d(num_channels, max_features, kernel_size=3, padding=1)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(max_features, num_channels, kernel_size=3, padding=1)
+    def forward(self, kp_source, kp_driving, jacobian):
+        x = self.relu(self.conv1(kp_source))
+        x = self.conv2(x)
+        sparse_motion = {'dense_motion': x}
+        return sparse_motion
+class Hourglass(nn.Module):
+    def __init__(self, block_expansion, num_blocks, max_features, num_channels, kp_size, num_deform_blocks):
+        super(Hourglass, self).__init__()
+        self.encoder = nn.Sequential(nn.Conv2d(num_channels, max_features, kernel_size=7, stride=2, padding=3),
+                                     nn.BatchNorm2d(max_features), nn.ReLU())
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(max_features, num_channels, kernel_size=4, stride=2, padding=1), nn.Tanh())
+    def forward(self, source_image, kp_driving, **kwargs):
+        x = self.encoder(source_image)
+        x = self.decoder(x)
+        B, C, H, W = x.size()
+        video = []
+        for _ in range(10):
+            frame = (x[0].cpu().detach().numpy().transpose(1, 2, 0) * 127.5 + 127.5).clip(0, 255).astype(
+                np.uint8)
+            video.append(frame)
+        return video
+class Face3DHelper:
+    def __init__(self, local_pca_path, device):
+        self.local_pca_path = local_pca_path
+        self.device = device
+    def run(self, source_image):
+        h, w, _ = source_image.shape
+        x_min = w // 4
+        y_min = h // 4
+        x_max = x_min + w // 2
+        y_max = y_min + h // 2
+        return [x_min, y_min, x_max, y_max]
+class Face3DHelperOld(Face3DHelper):
+    def __init__(self, local_pca_path, device):
+        super(Face3DHelperOld, self).__init__(local_pca_path, device)
+class MouthDetector:
+    def __init__(self):
+        pass
+    def detect(self, image):
+        h, w = image.shape[:2]
+        return (w // 2, h // 2)
+class KeypointNorm(nn.Module):
+    def __init__(self, device):
+        super(KeypointNorm, self).__init__()
+        self.device = device
+    def forward(self, kp_driving):
+        return kp_driving
+def save_video_with_watermark(video_frames, audio_path, output_path):
+    H, W, _ = video_frames[0].shape
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 25, (W, H))
+    for frame in video_frames:
+        out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    out.release()
+def paste_pic(video_path, source_image_crop, crop_info, audio_path, output_path):
+    shutil.copy(video_path, output_path)
+class TTSTalker:
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.tts_model = None
+    def load_model(self):
+        self.tts_model = self
+    def tokenizer(self, text):
+        return [ord(c) for c in text]
+    def __call__(self, input_tokens):
+        return torch.zeros(1, 16000, device=self.device)
+    def test(self, text, lang='en'):
+        if self.tts_model is None:
+            self.load_model()
+        output_path = os.path.join('./results', str(uuid.uuid4()) + '.wav')
+        os.makedirs('./results', exist_ok=True)
+        tokens = self.tokenizer(text)
+        input_tokens = torch.tensor([tokens], dtype=torch.long).to(self.device)
+        with torch.no_grad():
+            audio_output = self(input_tokens)
+        torchaudio.save(output_path, audio_output.cpu(), 16000)
+        return output_path
+class SadTalker:
+    def __init__(self, checkpoint_path='checkpoints', config_path='src/config', size=256, preprocess='crop',
+                 old_version=False):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.cfg = self.get_cfg_defaults()
+        self.merge_from_file(os.path.join(config_path, 'sadtalker_config.yaml'))
+        self.cfg['MODEL']['CHECKPOINTS_DIR'] = checkpoint_path
+        self.cfg['MODEL']['CONFIG_DIR'] = config_path
+        self.cfg['MODEL']['DEVICE'] = self.device
+        self.cfg['INPUT_IMAGE'] = {}
+        self.cfg['INPUT_IMAGE']['SOURCE_IMAGE'] = 'None'
+        self.cfg['INPUT_IMAGE']['DRIVEN_AUDIO'] = 'None'
+        self.cfg['INPUT_IMAGE']['PREPROCESS'] = preprocess
+        self.cfg['INPUT_IMAGE']['SIZE'] = size
+        self.cfg['INPUT_IMAGE']['OLD_VERSION'] = old_version
+        download_model(kp_url, kp_file, checkpoint_path)
+        download_model(aud_url, aud_file, checkpoint_path)
+        download_model(wav_url, wav_file, checkpoint_path)
+        download_model(gen_url, gen_file, checkpoint_path)
+        download_model(mapx_url, mapx_file, checkpoint_path)
+        download_model(den_url, den_file, checkpoint_path)
+        download_model(GFPGAN_URL, 'GFPGANv1.4.pth', checkpoint_path)
+        download_model(REALESRGAN_URL, 'RealESRGAN_x2plus.pth', checkpoint_path)
+        self.sadtalker_model = SadTalkerModel(self.cfg, device_id=[0])
+    def get_cfg_defaults(self):
+        return {
+            'MODEL': {
+                'CHECKPOINTS_DIR': '',
+                'CONFIG_DIR': '',
+                'DEVICE': self.device,
+                'SCALE': 64,
+                'NUM_VOXEL_FRAMES': 8,
+                'NUM_MOTION_FRAMES': 10,
+                'MAX_FEATURES': 256,
+                'DRIVEN_AUDIO_SAMPLE_RATE': 16000,
+                'VIDEO_FPS': 25,
+                'OUTPUT_VIDEO_FPS': None,
+                'OUTPUT_AUDIO_SAMPLE_RATE': None,
+                'USE_ENHANCER': False,
+                'ENHANCER_NAME': '',
+                'BG_UPSAMPLER': None,
+                'IS_HALF': False
+            },
+            'INPUT_IMAGE': {}
+        }
+    def merge_from_file(self, filepath):
+        if os.path.exists(filepath):
+            with open(filepath, 'r') as f:
+                cfg_from_file = yaml.safe_load(f)
+            self.cfg.update(cfg_from_file)
+    def test(self, source_image, driven_audio, preprocess='crop', still_mode=False, use_enhancer=False,
+             batch_size=1, size=256, pose_style=0, exp_scale=1.0, use_ref_video=False, ref_video=None,
+             ref_info=None, use_idle_mode=False, length_of_audio=0, use_blink=True, result_dir='./results/',
+             tts_text=None, tts_lang='en'):
+        self.sadtalker_model.test(source_image, driven_audio, preprocess, still_mode, use_enhancer, batch_size, size,
+                                  pose_style, exp_scale, use_ref_video, ref_video, ref_info, use_idle_mode,
+                                  length_of_audio, use_blink, result_dir, tts_text, tts_lang)
+        return self.sadtalker_model.save_result()
+class SadTalkerModel:
+    def __init__(self, sadtalker_cfg, device_id=[0]):
+        self.cfg = sadtalker_cfg
+        self.device = sadtalker_cfg['MODEL'].get('DEVICE', 'cpu')
+        self.sadtalker = SadTalkerInnerModel(sadtalker_cfg, device_id)
+        self.preprocesser = self.sadtalker.preprocesser
+        self.kp_extractor = self.sadtalker.kp_extractor
+        self.generator = self.sadtalker.generator
+        self.mapping = self.sadtalker.mapping
+        self.he_estimator = self.sadtalker.he_estimator
+        self.audio_to_coeff = self.sadtalker.audio_to_coeff
+        self.animate_from_coeff = self.sadtalker.animate_from_coeff
+        self.face_enhancer = self.sadtalker.face_enhancer
+    def test(self, source_image, driven_audio, preprocess='crop', still_mode=False, use_enhancer=False,
+             batch_size=1, size=256, pose_style=0, exp_scale=1.0, use_ref_video=False, ref_video=None,
+             ref_info=None, use_idle_mode=False, length_of_audio=0, use_blink=True, result_dir='./results/',
+             tts_text=None, tts_lang='en', jitter_amount=10, jitter_source_image=False):
+        self.inner_test = SadTalkerInner(self, source_image, driven_audio, preprocess, still_mode, use_enhancer,
+                                        batch_size, size, pose_style, exp_scale, use_ref_video, ref_video, ref_info,
+                                        use_idle_mode, length_of_audio, use_blink, result_dir, tts_text, tts_lang,
+                                        jitter_amount, jitter_source_image)
+        return self.inner_test.test()
+    def save_result(self):
+        return self.inner_test.save_result()
+class SadTalkerInner:
+    def __init__(self, sadtalker_model, source_image, driven_audio, preprocess, still_mode, use_enhancer,
+                 batch_size, size, pose_style, exp_scale, use_ref_video, ref_video, ref_info, use_idle_mode,
+                 length_of_audio, use_blink, result_dir, tts_text, tts_lang, jitter_amount, jitter_source_image):
+        self.sadtalker_model = sadtalker_model
+        self.source_image = source_image
+        self.driven_audio = driven_audio
+        self.preprocess = preprocess
+        self.still_mode = still_mode
+        self.use_enhancer = use_enhancer
+        self.batch_size = batch_size
+        self.size = size
+        self.pose_style = pose_style
+        self.exp_scale = exp_scale
+        self.use_ref_video = use_ref_video
+        self.ref_video = ref_video
+        self.ref_info = ref_info
+        self.use_idle_mode = use_idle_mode
+        self.length_of_audio = length_of_audio
+        self.use_blink = use_blink
+        self.result_dir = result_dir
+        self.tts_text = tts_text
+        self.tts_lang = tts_lang
+        self.jitter_amount = jitter_amount
+        self.jitter_source_image = jitter_source_image
+        self.device = self.sadtalker_model.device
+        self.output_path = None
+    def get_test_data(self):
+        proc = self.sadtalker_model.preprocesser
+        if self.tts_text is not None:
+            temp_dir = tempfile.mkdtemp()
+            audio_path = os.path.join(temp_dir, 'audio.wav')
+            tts = TTSTalker()
+            tts.test(self.tts_text, self.tts_lang)
+            self.driven_audio = audio_path
+        source_image_pil = Image.open(self.source_image).convert('RGB')
+        if self.jitter_source_image:
+            jitter_dx = np.random.randint(-self.jitter_amount, self.jitter_amount + 1)
+            jitter_dy = np.random.randint(-self.jitter_amount, self.jitter_amount + 1)
+            source_image_pil = Image.fromarray(
+                np.roll(np.roll(np.array(source_image_pil), jitter_dx, axis=1), jitter_dy, axis=0))
+        source_image_tensor, crop_info, cropped_image = proc.crop(source_image_pil, self.preprocess, self.size)
+        if self.still_mode or self.use_idle_mode:
+            ref_pose_coeff = proc.generate_still_pose(self.pose_style)
+            ref_expression_coeff = proc.generate_still_expression(self.exp_scale)
+        elif self.use_idle_mode:
+            ref_pose_coeff = proc.generate_idles_pose(self.length_of_audio, self.pose_style)
+            ref_expression_coeff = proc.generate_idles_expression(self.length_of_audio)
+        else:
+            ref_pose_coeff = None
+            ref_expression_coeff = None
+        audio_tensor, audio_sample_rate = proc.process_audio(self.driven_audio,
+                                                             self.sadtalker_model.cfg['MODEL']['DRIVEN_AUDIO_SAMPLE_RATE'])
+        batch = {
+            'source_image': source_image_tensor.unsqueeze(0).to(self.device),
+            'audio': audio_tensor.unsqueeze(0).to(self.device),
+            'ref_pose_coeff': ref_pose_coeff,
+            'ref_expression_coeff': ref_expression_coeff,
+            'source_image_crop': cropped_image,
+            'crop_info': crop_info,
+            'use_blink': self.use_blink,
+            'pose_style': self.pose_style,
+            'exp_scale': self.exp_scale,
+            'ref_video': self.ref_video,
+            'use_ref_video': self.use_ref_video,
+            'ref_info': self.ref_info,
+        }
+        return batch, audio_sample_rate
+    def run_inference(self, batch):
+        kp_extractor = self.sadtalker_model.kp_extractor
+        generator = self.sadtalker_model.generator
+        mapping = self.sadtalker_model.mapping
+        he_estimator = self.sadtalker_model.he_estimator
+        audio_to_coeff = self.sadtalker_model.audio_to_coeff
+        animate_from_coeff = self.sadtalker_model.animate_from_coeff
+        proc = self.sadtalker_model.preprocesser
+        with torch.no_grad():
+            kp_source = kp_extractor(batch['source_image'])
+            if self.still_mode or self.use_idle_mode:
+                ref_pose_coeff = batch['ref_pose_coeff']
+                ref_expression_coeff = batch['ref_expression_coeff']
+                pose_coeff = audio_to_coeff.get_pose_coeff(batch['audio'], ref_pose_coeff)
+                expression_coeff = audio_to_coeff.get_exp_coeff(batch['audio'], ref_expression_coeff)
+            elif self.use_idle_mode:
+                ref_pose_coeff = batch['ref_pose_coeff']
+                ref_expression_coeff = batch['ref_expression_coeff']
+                pose_coeff = audio_to_coeff.get_pose_coeff(batch['audio'], ref_pose_coeff)
+                expression_coeff = audio_to_coeff.get_exp_coeff(batch['audio'], ref_expression_coeff)
+            else:
+                if self.use_ref_video:
+                    kp_ref = kp_extractor(batch['source_image'])
+                    pose_coeff = audio_to_coeff.get_pose_coeff(batch['audio'], kp_ref=kp_ref,
+                                                                use_ref_info=batch['ref_info'])
+                else:
+                    pose_coeff = audio_to_coeff.get_pose_coeff(batch['audio'])
+                expression_coeff = audio_to_coeff.get_exp_coeff(batch['audio'])
+            coeff = {'pose_coeff': pose_coeff, 'expression_coeff': expression_coeff}
+            if self.use_blink:
+                coeff['blink_coeff'] = audio_to_coeff.get_blink_coeff(batch['audio'])
+            else:
+                coeff['blink_coeff'] = None
+            kp_driving = audio_to_coeff(batch['audio'])[0]
+            kp_norm = animate_from_coeff.normalize_kp(kp_driving)
+            coeff['kp_driving'] = kp_norm
+            coeff['jacobian'] = [torch.eye(2).unsqueeze(0).unsqueeze(0).to(self.device)] * 4
+            face_enhancer = self.sadtalker_model.face_enhancer if self.use_enhancer else None
+            output_video = animate_from_coeff.generate(batch['source_image'], kp_source, coeff, generator, mapping,
+                                                        he_estimator, batch['audio'], batch['source_image_crop'],
+                                                        face_enhancer=face_enhancer)
+        return output_video
+    def post_processing(self, output_video, audio_sample_rate, batch):
+        proc = self.sadtalker_model.preprocesser
+        base_name = os.path.splitext(os.path.basename(batch['source_image_crop']))[0]
+        audio_name = os.path.splitext(os.path.basename(self.driven_audio))[0]
+        output_video_path = os.path.join(self.result_dir, base_name + '_' + audio_name + '.mp4')
+        self.output_path = output_video_path
+        video_fps = self.sadtalker_model.cfg['MODEL']['VIDEO_FPS'] if self.sadtalker_model.cfg['MODEL'][
+                                                                          'OUTPUT_VIDEO_FPS'] is None else \
+            self.sadtalker_model.cfg['MODEL']['OUTPUT_VIDEO_FPS']
+        audio_output_sample_rate = self.sadtalker_model.cfg['MODEL']['DRIVEN_AUDIO_SAMPLE_RATE'] if \
+            self.sadtalker_model.cfg['MODEL']['OUTPUT_AUDIO_SAMPLE_RATE'] is None else \
+            self.sadtalker_model.cfg['MODEL']['OUTPUT_AUDIO_SAMPLE_RATE']
+        if self.use_enhancer:
+            enhanced_path = os.path.join(self.result_dir, base_name + '_' + audio_name + '_enhanced.mp4')
+            save_video_with_watermark(output_video, self.driven_audio, enhanced_path)
+            paste_pic(enhanced_path, batch['source_image_crop'], batch['crop_info'], self.driven_audio,
+                      output_video_path)
+            os.remove(enhanced_path)
+        else:
+            save_video_with_watermark(output_video, self.driven_audio, output_video_path)
+        if self.tts_text is not None:
+            shutil.rmtree(os.path.dirname(self.driven_audio))
+    def save_result(self):
+        return self.output_path
+    def __call__(self):
+        return self.output_path
+    def test(self):
+        batch, audio_sample_rate = self.get_test_data()
+        output_video = self.run_inference(batch)
+        self.post_processing(output_video, audio_sample_rate, batch)
+        return self.save_result()
+class SadTalkerInnerModel:
+    def __init__(self, sadtalker_cfg, device_id=[0]):
+        self.cfg = sadtalker_cfg
+        self.device = sadtalker_cfg['MODEL'].get('DEVICE', 'cpu')
+        self.preprocesser = Preprocesser(sadtalker_cfg, self.device)
+        self.kp_extractor = KeyPointExtractor(sadtalker_cfg, self.device)
+        self.audio_to_coeff = Audio2Coeff(sadtalker_cfg, self.device)
+        self.animate_from_coeff = AnimateFromCoeff(sadtalker_cfg, self.device)
+        self.face_enhancer = FaceEnhancer(sadtalker_cfg, self.device) if sadtalker_cfg['MODEL'][
+            'USE_ENHANCER'] else None
+        self.generator = Generator(sadtalker_cfg, self.device)
+        self.mapping = Mapping(sadtalker_cfg, self.device)
+        self.he_estimator = OcclusionAwareDenseMotion(sadtalker_cfg, self.device)
+class Preprocesser:
+    def __init__(self, sadtalker_cfg, device):
+        self.cfg = sadtalker_cfg
+        self.device = device
+        if self.cfg['INPUT_IMAGE'].get('OLD_VERSION', False):
+            self.face3d_helper = Face3DHelperOld(self.cfg['INPUT_IMAGE'].get('LOCAL_PCA_PATH', ''), device)
+        else:
+            self.face3d_helper = Face3DHelper(self.cfg['INPUT_IMAGE'].get('LOCAL_PCA_PATH', ''), device)
+        self.mouth_detector = MouthDetector()
+    def crop(self, source_image_pil, preprocess_type, size=256):
+        source_image = np.array(source_image_pil)
+        face_info = self.face3d_helper.run(source_image)
+        if face_info is None:
+            raise Exception("No face detected")
+        x_min, y_min, x_max, y_max = face_info[:4]
+        old_size = (x_max - x_min, y_max - y_min)
+        x_center = (x_max + x_min) / 2
+        y_center = (y_max + y_min) / 2
+        if preprocess_type == 'crop':
+            face_size = max(x_max - x_min, y_max - y_min)
+            x_min = int(x_center - face_size / 2)
+            y_min = int(y_center - face_size / 2)
+            x_max = int(x_center + face_size / 2)
+            y_max = int(y_center + face_size / 2)
+        else:
+            x_min -= int((x_max - x_min) * 0.1)
+            y_min -= int((y_max - y_min) * 0.1)
+            x_max += int((x_max - x_min) * 0.1)
+            y_max += int((y_max - y_min) * 0.1)
+        h, w = source_image.shape[:2]
+        x_min = max(0, x_min)
+        y_min = max(0, y_min)
+        x_max = min(w, x_max)
+        y_max = min(h, y_max)
+        cropped_image = source_image[y_min:y_max, x_min:x_max]
+        cropped_image_pil = Image.fromarray(cropped_image)
+        if size is not None and size != 0:
+            cropped_image_pil = cropped_image_pil.resize((size, size), Image.Resampling.LANCZOS)
+        source_image_tensor = self.img2tensor(cropped_image_pil)
+        return source_image_tensor, [[y_min, y_max], [x_min, x_max], old_size, cropped_image_pil.size], os.path.basename(
+            self.cfg['INPUT_IMAGE'].get('SOURCE_IMAGE', ''))
+    def img2tensor(self, img):
+        img = np.array(img).astype(np.float32) / 255.0
+        img = np.transpose(img, (2, 0, 1))
+        return torch.FloatTensor(img)
+    def video_to_tensor(self, video, device):
+        video_tensor_list = []
+        import torchvision.transforms as transforms
+        transform_func = transforms.ToTensor()
+        for frame in video:
+            frame_pil = Image.fromarray(frame)
+            frame_tensor = transform_func(frame_pil).unsqueeze(0).to(device)
+            video_tensor_list.append(frame_tensor)
+        video_tensor = torch.cat(video_tensor_list, dim=0)
+        return video_tensor
+    def process_audio(self, audio_path, sample_rate):
+        wav = load_wav_util(audio_path, sample_rate)
+        wav_tensor = torch.FloatTensor(wav).unsqueeze(0)
+        return wav_tensor, sample_rate
+    def generate_still_pose(self, pose_style):
+        ref_pose_coeff = torch.zeros((1, 64), dtype=torch.float32).to(self.device)
+        ref_pose_coeff[:, :3] = torch.tensor([0, 0, pose_style * 0.3], dtype=torch.float32)
+        return ref_pose_coeff
+    def generate_still_expression(self, exp_scale):
+        ref_expression_coeff = torch.zeros((1, 64), dtype=torch.float32).to(self.device)
+        ref_expression_coeff[:, :3] = torch.tensor([0, 0, exp_scale * 0.3], dtype=torch.float32)
+        return ref_expression_coeff
+    def generate_idles_pose(self, length_of_audio, pose_style):
+        num_frames = int(length_of_audio * self.cfg['MODEL']['VIDEO_FPS'])
+        ref_pose_coeff = torch.zeros((num_frames, 64), dtype=torch.float32).to(self.device)
+        start_pose = self.generate_still_pose(pose_style)
+        end_pose = self.generate_still_pose(pose_style)
+        for frame_idx in range(num_frames):
+            alpha = frame_idx / num_frames
+            ref_pose_coeff[frame_idx] = (1 - alpha) * start_pose + alpha * end_pose
+        return ref_pose_coeff
+    def generate_idles_expression(self, length_of_audio):
+        num_frames = int(length_of_audio * self.cfg['MODEL']['VIDEO_FPS'])
+        ref_expression_coeff = torch.zeros((num_frames, 64), dtype=torch.float32).to(self.device)
+        start_exp = self.generate_still_expression(1.0)
+        end_exp = self.generate_still_expression(1.0)
+        for frame_idx in range(num_frames):
+            alpha = frame_idx / num_frames
+            ref_expression_coeff[frame_idx] = (1 - alpha) * start_exp + alpha * end_exp
+        return ref_expression_coeff
+class KeyPointExtractor(nn.Module):
+    def __init__(self, sadtalker_cfg, device):
+        super(KeyPointExtractor, self).__init__()
+        self.kp_extractor = OcclusionAwareKPDetector(kp_channels=sadtalker_cfg['MODEL']['NUM_MOTION_FRAMES'],
+                                                      num_kp=10,
+                                                      num_dilation_blocks=2,
+                                                      dropout_rate=0.1).to(device)
+        checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'kp_detector.safetensors')
+        self.load_kp_detector(checkpoint_path, device)
+    def load_kp_detector(self, checkpoint_path, device):
+        if os.path.exists(checkpoint_path):
+            if checkpoint_path.endswith('safetensors'):
+                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location=device)
+            self.kp_extractor.load_state_dict(checkpoint.get('kp_detector', {}))
+        else:
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    def forward(self, x):
+        kp = self.kp_extractor(x)
+        return kp
+class Audio2Coeff(nn.Module):
+    def __init__(self, sadtalker_cfg, device):
+        super(Audio2Coeff, self).__init__()
+        self.audio_model = Wav2Vec2Model().to(device)
+        checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'wav2vec2.pth')
+        self.load_audio_model(checkpoint_path, device)
+        self.pose_mapper = AudioCoeffsPredictor(2048, 64).to(device)
+        self.exp_mapper = AudioCoeffsPredictor(2048, 64).to(device)
+        self.blink_mapper = AudioCoeffsPredictor(2048, 1).to(device)
+        mapping_checkpoint = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'audio2pose_00140-model.pth')
+        self.load_mapping_model(mapping_checkpoint, device)
+    def load_audio_model(self, checkpoint_path, device):
+        if os.path.exists(checkpoint_path):
+            if checkpoint_path.endswith('safetensors'):
+                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location=device)
+            self.audio_model.load_state_dict(checkpoint.get("wav2vec2", {}))
+        else:
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    def load_mapping_model(self, checkpoint_path, device):
+        if os.path.exists(checkpoint_path):
+            if checkpoint_path.endswith('safetensors'):
+                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location=device)
+            self.pose_mapper.load_state_dict(checkpoint.get("pose_predictor", {}))
+            self.exp_mapper.load_state_dict(checkpoint.get("exp_predictor", {}))
+            self.blink_mapper.load_state_dict(checkpoint.get("blink_predictor", {}))
+        else:
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    def get_pose_coeff(self, audio_tensor, ref_pose_coeff=None, kp_ref=None, use_ref_info=''):
+        audio_embedding = self.audio_model(audio_tensor)
+        pose_coeff = self.pose_mapper(audio_embedding)
+        if ref_pose_coeff is not None:
+            pose_coeff = ref_pose_coeff
+        if kp_ref is not None and use_ref_info == 'pose':
+            ref_pose_6d = kp_ref['value'][:, :6]
+            pose_coeff[:, :6] = self.mean_std_normalize(ref_pose_6d).mean(dim=1)
+        return pose_coeff
+    def get_exp_coeff(self, audio_tensor, ref_expression_coeff=None):
+        audio_embedding = self.audio_model(audio_tensor)
+        expression_coeff = self.exp_mapper(audio_embedding)
+        if ref_expression_coeff is not None:
+            expression_coeff = ref_expression_coeff
+        return expression_coeff
+    def get_blink_coeff(self, audio_tensor):
+        audio_embedding = self.audio_model(audio_tensor)
+        blink_coeff = self.blink_mapper(audio_embedding)
+        return blink_coeff
+    def forward(self, audio):
+        audio_embedding = self.audio_model(audio)
+        pose_coeff, expression_coeff, blink_coeff = self.pose_mapper(audio_embedding), self.exp_mapper(
+            audio_embedding), self.blink_mapper(audio_embedding)
+        return pose_coeff, expression_coeff, blink_coeff
+    def mean_std_normalize(self, coeff):
+        mean = coeff.mean(dim=1, keepdim=True)
+        std = coeff.std(dim=1, keepdim=True)
+        return (coeff - mean) / std
+class AnimateFromCoeff(nn.Module):
+    def __init__(self, sadtalker_cfg, device):
+        super(AnimateFromCoeff, self).__init__()
+        self.generator = Generator(sadtalker_cfg, device)
+        self.mapping = Mapping(sadtalker_cfg, device)
+        self.kp_norm = KeypointNorm(device=device)
+        self.he_estimator = OcclusionAwareDenseMotion(sadtalker_cfg, device)
+    def normalize_kp(self, kp_driving):
+        return self.kp_norm(kp_driving)
+    def generate(self, source_image, kp_source, coeff, generator, mapping, he_estimator, audio, source_image_crop,
+                 face_enhancer=None):
+        kp_driving = coeff['kp_driving']
+        jacobian = coeff['jacobian']
+        pose_coeff = coeff['pose_coeff']
+        expression_coeff = coeff['expression_coeff']
+        blink_coeff = coeff['blink_coeff']
+        with torch.no_grad():
+            if blink_coeff is not None:
+                sparse_motion = he_estimator(kp_source, kp_driving, jacobian)
+                dense_motion = sparse_motion['dense_motion']
+                video_deocclusion = generator(source_image, dense_motion, bg_param={'mask': None, 'color': None})
+                face_3d = mapping(expression_coeff, pose_coeff, blink_coeff)
+                video_3d = generator(source_image, dense_motion, bg_param={'mask': None, 'color': None},
+                                      face_3d_param=face_3d)
+                video_output = video_deocclusion['video_no_reocclusion'] + video_3d['video_3d']
+                video_output = self.make_animation(video_output)
+            else:
+                sparse_motion = he_estimator(kp_source, kp_driving, jacobian)
+                dense_motion = sparse_motion['dense_motion']
+                face_3d = mapping(expression_coeff, pose_coeff)
+                video_3d = generator(source_image, dense_motion, bg_param={'mask': None, 'color': None},
+                                      face_3d_param=face_3d)
+                video_output = video_3d['video_3d']
+                video_output = self.make_animation(video_output)
+            if face_enhancer is not None:
+                video_output_enhanced = []
+                for frame in tqdm(video_output, 'Face enhancer running'):
+                    pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+                    enhanced_image = face_enhancer.enhance(np.array(pil_image))[0]
+                    video_output_enhanced.append(cv2.cvtColor(enhanced_image, cv2.COLOR_BGR2RGB))
+                video_output = video_output_enhanced
+        return video_output
+    def make_animation(self, video_array):
+        H, W, _ = video_array[0].shape
+        out = cv2.VideoWriter('./tmp.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 25, (W, H))
+        for img in video_array:
+            out.write(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+        out.release()
+        video = imageio.mimread('./tmp.mp4')
+        os.remove('./tmp.mp4')
+        return video
+class Generator(nn.Module):
+    def __init__(self, sadtalker_cfg, device):
+        super(Generator, self).__init__()
+        self.generator = Hourglass(block_expansion=sadtalker_cfg['MODEL']['SCALE'],
+                                     num_blocks=sadtalker_cfg['MODEL']['NUM_VOXEL_FRAMES'],
+                                     max_features=sadtalker_cfg['MODEL']['MAX_FEATURES'],
+                                     num_channels=3,
+                                     kp_size=10,
+                                     num_deform_blocks=sadtalker_cfg['MODEL']['NUM_MOTION_FRAMES']).to(device)
+        checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'generator.pth')
+        self.load_generator(checkpoint_path, device)
+    def load_generator(self, checkpoint_path, device):
+        if os.path.exists(checkpoint_path):
+            if checkpoint_path.endswith('safetensors'):
+                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location=device)
+            self.generator.load_state_dict(checkpoint.get('generator', {}))
+        else:
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    def forward(self, source_image, dense_motion, bg_param, face_3d_param=None):
+        if face_3d_param is not None:
+            video_3d = self.generator(source_image, kp_driving=dense_motion, bg_param=bg_param,
+                                      face_3d_param=face_3d_param)
+        else:
+            video_3d = self.generator(source_image, kp_driving=dense_motion, bg_param=bg_param)
+        return {'video_3d': video_3d, 'video_no_reocclusion': video_3d}
+class Mapping(nn.Module):
+    def __init__(self, sadtalker_cfg, device):
+        super(Mapping, self).__init__()
+        self.mapping_net = MappingNet(num_coeffs=64, num_layers=3, hidden_dim=128).to(device)
+        checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'mapping.pth')
+        self.load_mapping_net(checkpoint_path, device)
+        self.f_3d_mean = torch.zeros(1, 64, device=device)
+    def load_mapping_net(self, checkpoint_path, device):
+        if os.path.exists(checkpoint_path):
+            if checkpoint_path.endswith('safetensors'):
+                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location=device)
+            self.mapping_net.load_state_dict(checkpoint.get('mapping', {}))
+        else:
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    def forward(self, expression_coeff, pose_coeff, blink_coeff=None):
+        coeff = torch.cat([expression_coeff, pose_coeff], dim=1)
+        face_3d = self.mapping_net(coeff) + self.f_3d_mean
+        if blink_coeff is not None:
+            face_3d[:, -1:] = blink_coeff
+        return face_3d
+class OcclusionAwareDenseMotion(nn.Module):
+    def __init__(self, sadtalker_cfg, device):
+        super(OcclusionAwareDenseMotion, self).__init__()
+        self.dense_motion_network = DenseMotionNetwork(num_kp=10,
+                                                        num_channels=3,
+                                                        block_expansion=sadtalker_cfg['MODEL']['SCALE'],
+                                                        num_blocks=sadtalker_cfg['MODEL']['NUM_MOTION_FRAMES'] - 1,
+                                                        max_features=sadtalker_cfg['MODEL']['MAX_FEATURES']).to(device)
+        checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'dense_motion.pth')
+        self.load_dense_motion_network(checkpoint_path, device)
+    def load_dense_motion_network(self, checkpoint_path, device):
+        if os.path.exists(checkpoint_path):
+            if checkpoint_path.endswith('safetensors'):
+                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location=device)
+            self.dense_motion_network.load_state_dict(checkpoint.get('dense_motion', {}))
+        else:
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    def forward(self, kp_source, kp_driving, jacobian):
+        sparse_motion = self.dense_motion_network(kp_source, kp_driving, jacobian)
+        return sparse_motion
+class FaceEnhancer(nn.Module):
+    def __init__(self, sadtalker_cfg, device):
+        super(FaceEnhancer, self).__init__()
+        enhancer_name = sadtalker_cfg['MODEL']['ENHANCER_NAME']
+        bg_upsampler = sadtalker_cfg['MODEL']['BG_UPSAMPLER']
+        if enhancer_name == 'gfpgan':
+            from gfpgan import GFPGANer
+            self.face_enhancer = GFPGANer(model_path=os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'GFPGANv1.4.pth'),
+                                          upscale=1,
+                                          arch='clean',
+                                          channel_multiplier=2,
+                                          bg_upsampler=bg_upsampler)
+        elif enhancer_name == 'realesrgan':
+            from realesrgan import RealESRGANer
+            half = False if device == 'cpu' else sadtalker_cfg['MODEL']['IS_HALF']
+            self.face_enhancer = RealESRGANer(scale=2,
+                                               model_path=os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'],
+                                                                      'RealESRGAN_x2plus.pth'),
+                                               tile=0,
+                                               tile_pad=10,
+                                               pre_pad=0,
+                                               half=half,
+                                               device=device)
+        else:
+            self.face_enhancer = None
+    def forward(self, x):
+        return self.face_enhancer.enhance(x, outscale=1)[0]

sentiment_api.py ADDED Viewed

	@@ -0,0 +1,27 @@

+from flask import jsonify
+from main import *
+#from main import import sentiment_model, device
+import torch
+def analyze_sentiment(text, output_path="output_sentiment.json"):
+    if sentiment_model is None:
+        return "Sentiment model not initialized."
+    input_tokens = sentiment_model.tokenizer(text, return_tensors="pt", padding=True).to(device)
+    with torch.no_grad():
+        sentiment_logits = sentiment_model(input_tokens['input_ids'])
+    predicted_class_id = torch.argmax(sentiment_logits, dim=-1).item()
+    sentiment_label = sentiment_model.config.id2label[predicted_class_id]
+    probability = torch.softmax(sentiment_logits, dim=-1)[0][predicted_class_id].item()
+    return {"sentiment": sentiment_label, "probability": probability}
+def sentiment_api():
+    data = request.get_json()
+    text = data.get('text')
+    if not text:
+        return jsonify({"error": "Text is required"}), 400
+    output_file = analyze_sentiment(text)
+    if output_file == "Sentiment model not initialized.":
+        return jsonify({"error": "Sentiment analysis failed"}), 500
+    return jsonify(output_file)

stt_api.py ADDED Viewed

	@@ -0,0 +1,36 @@

+import os
+import uuid
+from flask import jsonify, send_file, request
+from main import *
+#from main import import stt_model, device
+import torch
+import torchaudio
+def speech_to_text_func(audio_path, output_path="output_stt.txt"):
+    if stt_model is None:
+        return "STT model not initialized."
+    waveform, sample_rate = torchaudio.load(audio_path)
+    if waveform.ndim > 1:
+        waveform = torch.mean(waveform, dim=0, keepdim=True)
+    waveform = waveform.to(device)
+    with torch.no_grad():
+        logits = stt_model(waveform)
+    predicted_ids = torch.argmax(logits, dim=-1)
+    transcription = stt_model.tokenizer.decode(predicted_ids[0].cpu().tolist())
+    with open(output_path, "w") as file:
+        file.write(transcription)
+    return output_path
+def stt_api():
+    if 'audio' not in request.files:
+        return jsonify({"error": "Audio file is required"}), 400
+    audio_file = request.files['audio']
+    temp_audio_path = f"temp_audio_{uuid.uuid4()}.wav"
+    audio_file.save(temp_audio_path)
+    output_file = speech_to_text_func(temp_audio_path)
+    os.remove(temp_audio_path)
+    if output_file == "STT model not initialized.":
+        return jsonify({"error": "STT failed"}), 500
+    return send_file(output_file, mimetype="text/plain", as_attachment=True, download_name="output.txt")

summarization_api.py ADDED Viewed

	@@ -0,0 +1,29 @@

+from flask import jsonify, send_file, request
+from main import *
+#from main import import summarization_model, summarization_word_to_index, device
+import torch
+def summarize_text(text, output_path="output_summary.txt"):
+    if summarization_model is None:
+        return "Summarization model not initialized."
+    input_tokens = [summarization_word_to_index.get(word.lower(), 1) for word in text.split()]
+    input_tensor = torch.tensor([input_tokens], dtype=torch.long).to(device)
+    with torch.no_grad():
+        summary_ids = summarization_model.generate(input_tensor, num_beams=4, max_length=100, early_stopping=True)
+        summary_text = summarization_model.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    with open(output_path, "w") as file:
+        file.write(summary_text)
+    return output_path
+def summarization_api():
+    data = request.get_json()
+    text = data.get('text')
+    if not text:
+        return jsonify({"error": "Text is required"}), 400
+    output_file = summarize_text(text)
+    if output_file == "Summarization model not initialized.":
+        return jsonify({"error": "Summarization failed"}), 500
+    return send_file(output_file, mimetype="text/plain", as_attachment=True, download_name="output_summary.txt")

text_generation.py ADDED Viewed

	@@ -0,0 +1,152 @@

+import torch
+import torch.nn.functional as F
+from tqdm import trange
+import time
+from tokenxxx import *
+from main import *
+#from main import import model_gpt2, enc, codegen_model, codegen_tokenizer, summarization_model, device, system_prompt, MAX_LENGTH, summarize_text as summarize_func
+from duckduckgo_search import DDGS
+def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
+    top_k = min(top_k, logits.size(-1))
+    if top_k > 0:
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., [-1]]
+        logits[indices_to_remove] = filter_value
+    if top_p > 0.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+        indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        logits[indices_to_remove] = filter_value
+    return logits
+def sample_sequence(prompt, model, enc, length, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"):
+    start_time = time.time()
+    context_tokens = enc.encode(prompt)
+    context_tokens_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device)
+    generated = context_tokens
+    past = None
+    text_generated_count = 0
+    past_key_values = past if past is not None else None
+    with torch.no_grad():
+        outputs = model(context_tokens_tensor, past_key_values=past_key_values)
+        next_token_logits = outputs[0][:, -1, :] / temperature
+        past = outputs[1]
+        for token_index in set(generated):
+            next_token_logits[0, token_index] /= repetition_penalty
+        filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)
+        if temperature == 0:
+            next_token = torch.argmax(filtered_logits, dim=-1).unsqueeze(0)
+        else:
+            next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)
+        generated += next_token.tolist()[0]
+        text_generated_count += 1
+        token = next_token.tolist()[0][0]
+        yield enc.decode([token])
+        if token == enc.encoder[END_OF_TEXT_TOKEN]:
+            yield "<END_STREAM>"
+        if text_generated_count > length:
+            yield "<END_STREAM>"
+        if (time.time() - start_time) * 1000 > 5000:
+            yield "<END_STREAM>"
+def sample_sequence_codegen(prompt, model, tokenizer, length, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"):
+    start_time = time.time()
+    context_tokens = tokenizer.encode(prompt)
+    context_tokens_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device).unsqueeze(0)
+    generated = context_tokens
+    past = None
+    text_generated_count = 0
+    with torch.no_grad():
+        outputs = model(input_ids=context_tokens_tensor, past_key_values=past, labels=None)
+        next_token_logits = outputs[0][:, -1, :] / temperature
+        past = outputs[1]
+        for token_index in set(generated):
+            next_token_logits[0, token_index] /= repetition_penalty
+        filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)
+        if temperature == 0:
+            next_token = torch.argmax(filtered_logits, dim=-1).unsqueeze(0)
+        else:
+            next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)
+        generated.append(next_token.tolist()[0][0])
+        text_generated_count += 1
+        token = next_token.tolist()[0][0]
+        yield tokenizer.decode([token])
+        if token == 50256:
+            yield "<END_STREAM>"
+        if text_generated_count > length:
+            yield "<END_STREAM>"
+        if (time.time() - start_time) * 1000 > 5000:
+            yield "<END_STREAM>"
+def perform_reasoning_stream(text_input, temperature, top_k, top_p, repetition_penalty):
+    try:
+        prompt_text = system_prompt + "\n\n"
+        prompt_text += "User: " + text_input + "\nCyrah: "
+        reasoning_prompt = prompt_text
+        ddgs = DDGS()
+        search_results = [r for r in ddgs.text(text_input, max_results=MAX_XDD)]
+        if search_results:
+            prompt_text += "\nWeb Search Results:\n"
+            for result in search_results:
+                prompt_text += f"- {result['body']}\n"
+            prompt_text += "\n"
+        generated_text_stream = []
+        stream_type = "text"
+        if "code" in text_input.lower() or "program" in text_input.lower():
+            if codegen_model and codegen_tokenizer:
+                generated_text_stream = sample_sequence_codegen(
+                    prompt=reasoning_prompt,
+                    model=codegen_model,
+                    tokenizer=codegen_tokenizer,
+                    length=MAX_LENGTH,
+                    temperature=temperature,
+                    top_k=top_k,
+                    top_p=top_p,
+                    repetition_penalty=repetition_penalty,
+                    device=device
+                )
+                stream_type = "text"
+        elif "summarize" in text_input.lower() or "summary" in text_input.lower():
+            if summarization_model:
+                summary = summarize_func(text_input)
+                yield f"SUMMARY_TEXT:{summary}"
+                yield "<END_STREAM>"
+                stream_type = "summary"
+        else:
+            if model_gpt2 and enc:
+                generated_text_stream = sample_sequence(
+                    prompt=reasoning_prompt,
+                    model=model_gpt2,
+                    enc=enc,
+                    length=MAX_LENGTH,
+                    temperature=temperature,
+                    top_k=top_k,
+                    top_p=top_p,
+                    repetition_penalty=repetition_penalty,
+                    device=device
+                )
+                stream_type = "text"
+        accumulated_text = ""
+        if stream_type == "text":
+            for token in generated_text_stream:
+                if token == "<END_STREAM>":
+                    yield accumulated_text
+                    yield "<END_STREAM>"
+                    return
+                if token == END_OF_TEXT_TOKEN:
+                    accumulated_text += END_OF_TEXT_TOKEN
+                    continue
+                if token:
+                    accumulated_text += token
+    except Exception as e:
+        print(f"Reasoning Error: {e}")
+        yield "Error during reasoning. Please try again."
+        yield "<END_STREAM>"

text_to_video_api.py ADDED Viewed

	@@ -0,0 +1,37 @@

+import os
+import uuid
+from flask import jsonify, send_file, request
+from main import *
+#from main import import text_to_video_model
+import torch
+import io
+from skimage import img_as_ubyte
+import imageio
+def text_to_video_func(prompt, output_path="output_video.mp4"):
+    if text_to_video_model is None:
+        return "Text-to-Video model not initialized."
+    video_frames_list = text_to_video_model(prompt)
+    if video_frames_list and hasattr(video_frames_list, 'frames'):
+        video_frames = video_frames_list.frames
+        export_to_video_pure(video_frames, output_video=output_path)
+        return output_path
+    return "Video generation failed."
+def export_to_video_pure(video_frames, output_video="output_video.mp4", fps=25):
+    writer = imageio.get_writer(output_video, fps=fps)
+    for frame in video_frames:
+        writer.append_data(img_as_ubyte(frame))
+    writer.close()
+def text_to_video_api():
+    data = request.get_json()
+    prompt = data.get('prompt')
+    if not prompt:
+        return jsonify({"error": "Prompt is required"}), 400
+    output_file = text_to_video_func(prompt)
+    if output_file == "Text-to-Video model not initialized." or output_file == "Video generation failed.":
+        return jsonify({"error": "Text to video failed"}), 500
+    with open(output_file, 'rb') as f:
+        video_content = f.read()
+    return send_file(io.BytesIO(video_content), mimetype='video/mp4', as_attachment=True, download_name="output_video.mp4")

tokenxxx.py ADDED Viewed

	@@ -0,0 +1,161 @@

+import json
+import re
+import unicodedata
+from functools import lru_cache
+import wget
+import os
+from constants import *
+import nltk
+@lru_cache()
+def bytes_to_unicode():
+    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+def get_pairs(word):
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+class Encoder:
+    def __init__(self, encoder, bpe_merges, errors='replace', tokenize=None):
+        self.encoder = encoder
+        self.decoder = {v:k for k,v in self.encoder.items()}
+        self.errors = errors
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v:k for k, v in self.byte_encoder.items()}
+        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
+        self.cache = {}
+        if tokenize is None:
+            self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\w+| ?[^\s\w]+|\s+(?!\S)|\s+""", re.UNICODE)
+            self.tokenize = lambda text: re.findall(self.pat, text)
+        else:
+            self.tokenize = tokenize
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token)
+        pairs = get_pairs(word)
+        if not pairs:
+            return token
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+    def encode(self, text):
+        bpe_tokens = []
+        normalized_text = unicodedata.normalize('NFKC', text)
+        normalized_text = ''.join(c for c in normalized_text if c.isascii() and c != '\t')
+        normalized_text = ''.join(c for c in normalized_text if not unicodedata.category(c).startswith('C'))
+        for token in self.tokenize(normalized_text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8', errors='ignore'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors='replace')
+        decoded_text = text.replace(" .", ".").replace(" ,", ",").replace(" '", "'").replace(" ?", "?").replace(" !", "!").replace(" :", ":").replace('\n', '<br>')
+        sentences = nltk.sent_tokenize(decoded_text)
+        return ' '.join(sentences).replace("<br>", "<br>\n")
+def get_encoder_gpt2():
+    encoder_path = os.path.join(GPT2_FOLDER, ENCODER_FILE)
+    vocab_path = os.path.join(GPT2_FOLDER, VOCAB_FILE)
+    if not os.path.exists(GPT2_FOLDER):
+        os.makedirs(GPT2_FOLDER)
+    if not os.path.exists(encoder_path):
+        wget.download(ENCODER_URL, out=encoder_path)
+    if not os.path.exists(vocab_path):
+        wget.download(VOCAB_URL, out=vocab_path)
+    with open(encoder_path, 'r') as f:
+        encoder = json.load(f)
+    with open(vocab_path, 'r', encoding="utf-8") as f:
+        bpe_data = f.read()
+    bpe_merges = [tuple(merge_str.split()) for merge_str in bpe_data.split('\n')[1:-1]]
+    encoder_obj = Encoder(encoder=encoder, bpe_merges=bpe_merges)
+    encoder_obj.encoder[END_OF_TEXT_TOKEN] = len(encoder_obj.encoder)
+    encoder_obj.decoder[len(encoder_obj.decoder)] = END_OF_TEXT_TOKEN
+    return encoder_obj
+def get_codegen_tokenizer_pure(vocab_file, merges_file):
+    vocab = json.load(open(vocab_file))
+    merges = open(merges_file, 'r', encoding="utf-8").read().split('\n')[1:-1]
+    bpe_merges = [tuple(m.split()) for m in merges]
+    byte_encoder = bytes_to_unicode()
+    byte_decoder = {v: k for k, v in byte_encoder.items()}
+    tokenizer_regex = re.compile(r'''<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+''')
+    tokenize = lambda text: re.findall(tokenizer_regex, text)
+    encoder_obj = Encoder(
+        encoder=vocab,
+        bpe_merges=bpe_merges,
+        byte_encoder=byte_encoder,
+        byte_decoder=byte_decoder,
+        tokenize=tokenize
+    )
+    return encoder_obj
+def codegen_tokenize(text, tokenizer):
+    return tokenizer.encode(text)
+def codegen_decode(tokens, tokenizer):
+    return tokenizer.decode(tokens)
+def tokenize_text(text):
+    global vocabulary, word_to_index, index_to_word
+    tokens = text.lower().split()
+    for token in tokens:
+        if token not in vocabulary:
+            vocabulary.add(token)
+            word_to_index[token] = len(index_to_word)
+            index_to_word.append(token)
+    return tokens
+def text_to_vector(text):
+    global vocabulary, word_to_index
+    tokens = tokenize_text(text)
+    vector = torch.zeros(len(vocabulary))
+    for token in tokens:
+        if token in word_to_index:
+            vector[word_to_index[token]] += 1
+    return vector

translation_api.py ADDED Viewed

	@@ -0,0 +1,27 @@

+from flask import jsonify, send_file, request
+from main import *
+#from main import import translation_model, device
+def perform_translation(text, target_language_code='es_XX', source_language_code='en_XX', output_path="output_translation.txt"):
+    if translation_model is None:
+        return "Translation model not initialized."
+    encoded_text = translation_model.tokenizer(text, return_tensors="pt", padding=True).to(device)
+    generated_tokens = translation_model.generate(input_ids=encoded_text['input_ids'], attention_mask=encoded_text['attention_mask'], forced_bos_token_id=translation_model.config.lang_code_to_id[target_language_code])
+    translation = translation_model.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    with open(output_path, "w") as file:
+        file.write(translation)
+    return output_path
+def translation_api():
+    data = request.get_json()
+    text = data.get('text')
+    target_lang = data.get('target_lang', 'es')
+    source_lang = data.get('source_lang', 'en')
+    if not text:
+        return jsonify({"error": "Text is required"}), 400
+    output_file = perform_translation(text, target_language_code=f'{target_lang}_XX', source_language_code=f'{source_lang}_XX')
+    if output_file == "Translation model not initialized.":
+        return jsonify({"error": "Translation failed"}), 500
+    return send_file(output_file, mimetype="text/plain", as_attachment=True, download_name="output_translation.txt")

tts_api.py ADDED Viewed

	@@ -0,0 +1,23 @@

+import os
+from flask import jsonify, send_file, request
+from main import *
+#from main import import tts_model, device
+def text_to_speech_func(text, output_path="output_tts.wav"):
+    if tts_model is None:
+        return "TTS model not initialized."
+    input_tokens = tts_model.tokenizer(text, return_tensors="pt", padding=True).to(device)
+    with torch.no_grad():
+        audio_output = tts_model(input_tokens['input_ids'])
+    torchaudio.save(output_path, audio_output.cpu(), 16000)
+    return output_path
+def tts_api():
+    data = request.get_json()
+    text = data.get('text')
+    if not text:
+        return jsonify({"error": "Text is required"}), 400
+    output_file = text_to_speech_func(text)
+    if output_file == "TTS model not initialized.":
+        return jsonify({"error": "TTS generation failed"}), 500
+    return send_file(output_file, mimetype="audio/wav", as_attachment=True, download_name="output.wav")

utils.py ADDED Viewed

	@@ -0,0 +1,190 @@

+import requests
+from bs4 import BeautifulSoup
+from faker import Faker
+from urllib.request import urlretrieve
+import urllib.request
+from urllib3.util.retry import Retry
+import time
+import os
+import wget
+import json
+import unicodedata
+import nltk
+from sklearn.datasets import fetch_20newsgroups
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.multiclass import OneVsRestClassifier
+import warnings
+from requests.adapters import HTTPAdapter
+from constants import  *
+MAX_XDD = 5
+use_google_search = True
+use_20newsgroup = True
+fake = Faker()
+def create_retry_session():
+    retry_strategy = Retry(
+        total=5,
+        status_forcelist=[429, 500, 502, 503, 504],
+        method_whitelist=["GET"],
+        backoff_factor=1,
+    )
+    adapter = HTTPAdapter(max_retries=retry_strategy)
+    http = requests.Session()
+    http.mount("https://", adapter)
+    http.mount("http://", adapter)
+    return http
+def get_google_search_results(query, retry_session):
+    if not use_google_search:
+        return []
+    headers = {"User-Agent": fake.user_agent()}
+    search_url = f"https://www.google.com/search?q={query}"
+    try:
+        response = retry_session.get(search_url, headers=headers, timeout=10)
+        response.raise_for_status()
+    except requests.exceptions.RequestException as e:
+        return []
+    soup = BeautifulSoup(response.text, "html.parser")
+    search_results = []
+    for a_tag in soup.find_all('a', href=True):
+        if 'url?q=' in a_tag['href'] and not a_tag['href'].startswith("https://accounts.google.com"):
+            search_results.append(a_tag['href'].split('url?q=')[1].split('&')[0])
+    return search_results
+def fetch_20newsgroup_data():
+    if not use_20newsgroup:
+        return []
+    try:
+        newsgroups_train = fetch_20newsgroups(subset='train', categories=['talk.trivia', 'rec.sport.baseball', 'sci.med', 'comp.sys.ibm.pc.hardware', 'soc.religion.christian'])
+        data = newsgroups_train.data
+        return data
+    except Exception as e:
+        return []
+def download_file(url, filename, folder, retries=3):
+    filepath = os.path.join(folder, filename)
+    if os.path.exists(filepath):
+        return True
+    os.makedirs(folder, exist_ok=True)
+    for attempt in range(retries):
+        try:
+            wget.download(url, out=filepath)
+            return True
+        except Exception as e:
+            if attempt < retries - 1:
+                time.sleep(2)
+            else:
+                return False
+    return False
+def download_gpt2_files(folder, model_url, model_file, encoder_url, encoder_file, vocab_url, vocab_file):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    if not os.path.exists(os.path.join(folder, model_file)):
+        download_file(model_url, model_file, folder)
+    if not os.path.exists(os.path.join(folder, encoder_file)):
+        download_file(encoder_url, encoder_file, folder)
+    if not os.path.exists(os.path.join(folder, vocab_file)):
+        download_file(vocab_url, vocab_file, folder)
+def download_translation_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_codegen_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_summarization_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_imagegen_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_image_to_3d_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_text_to_video_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_sentiment_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_stt_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_tts_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def download_musicgen_files(folder, model_files_urls):
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+    for url, filename in model_files_urls:
+        if not os.path.exists(os.path.join(folder, filename)):
+            download_file(url, filename, folder)
+def bytes_to_unicode_gpt2():
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n = n+1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+def get_codegen_tokenizer_pure(vocab_file, merges_file):
+    vocab = json.load(open(vocab_file))
+    merges = open(merges_file, 'r', encoding="utf-8").read().split('\n')[1:-1]
+    bpe_ranks = dict(zip(merges, range(len(merges))))
+    byte_encoder = bytes_to_unicode()
+    byte_decoder = {v: k for k, v in byte_encoder.items()}
+    tokenizer_regex = re.compile(r'''<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+''')
+    tokenize = lambda text: re.findall(tokenizer_regex, text)
+    encoder_obj = Encoder(
+        encoder=vocab,
+        decoder={v: u for u, v in vocab.items()},
+        bpe_ranks=bpe_ranks,
+        byte_encoder=byte_encoder,
+        byte_decoder=byte_decoder,
+        tokenize=tokenize
+    )
+    return encoder_obj

xxx.py ADDED Viewed

	@@ -0,0 +1,142 @@

+import json
+import re
+import unicodedata
+from functools import lru_cache
+import wget
+import os
+from constants import GPT2_FOLDER, ENCODER_FILE, VOCAB_FILE, END_OF_TEXT_TOKEN
+import nltk
+@lru_cache()
+def bytes_to_unicode():
+    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+def get_pairs(word):
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+class Encoder:
+    def __init__(self, encoder, bpe_merges, errors='replace', tokenize=None):
+        self.encoder = encoder
+        self.decoder = {v:k for k,v in self.encoder.items()}
+        self.errors = errors
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v:k for k, v in self.byte_encoder.items()}
+        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
+        self.cache = {}
+        if tokenize is None:
+            self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\w+| ?[^\s\w]+|\s+(?!\S)|\s+""", re.UNICODE)
+            self.tokenize = lambda text: re.findall(self.pat, text)
+        else:
+            self.tokenize = tokenize
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token)
+        pairs = get_pairs(word)
+        if not pairs:
+            return token
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+    def encode(self, text):
+        bpe_tokens = []
+        normalized_text = unicodedata.normalize('NFKC', text)
+        normalized_text = ''.join(c for c in normalized_text if c.isascii() and c != '\t')
+        normalized_text = ''.join(c for c in normalized_text if not unicodedata.category(c).startswith('C'))
+        for token in self.tokenize(normalized_text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8', errors='ignore'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors='replace')
+        decoded_text = text.replace(" .", ".").replace(" ,", ",").replace(" '", "'").replace(" ?", "?").replace(" !", "!").replace(" :", ":").replace('\n', '<br>')
+        sentences = nltk.sent_tokenize(decoded_text)
+        return ' '.join(sentences).replace("<br>", "<br>\n")
+def get_encoder_gpt2():
+    encoder_path = os.path.join(GPT2_FOLDER, ENCODER_FILE)
+    vocab_path = os.path.join(GPT2_FOLDER, VOCAB_FILE)
+    if not os.path.exists(GPT2_FOLDER):
+        os.makedirs(GPT2_FOLDER)
+    if not os.path.exists(encoder_path):
+        wget.download(ENCODER_URL, out=encoder_path)
+    if not os.path.exists(vocab_path):
+        wget.download(VOCAB_URL, out=vocab_path)
+    with open(encoder_path, 'r') as f:
+        encoder = json.load(f)
+    with open(vocab_path, 'r', encoding="utf-8") as f:
+        bpe_data = f.read()
+    bpe_merges = [tuple(merge_str.split()) for merge_str in bpe_data.split('\n')[1:-1]]
+    encoder_obj = Encoder(encoder=encoder, bpe_merges=bpe_merges)
+    encoder_obj.encoder[END_OF_TEXT_TOKEN] = len(encoder_obj.encoder)
+    encoder_obj.decoder[len(encoder_obj.decoder)] = END_OF_TEXT_TOKEN
+    return encoder_obj
+def get_codegen_tokenizer_pure(vocab_file, merges_file):
+    vocab = json.load(open(vocab_file))
+    merges = open(merges_file, 'r', encoding="utf-8").read().split('\n')[1:-1]
+    bpe_merges = [tuple(m.split()) for m in merges]
+    byte_encoder = bytes_to_unicode()
+    byte_decoder = {v: k for k, v in byte_encoder.items()}
+    tokenizer_regex = re.compile(r'''<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+''')
+    tokenize = lambda text: re.findall(tokenizer_regex, text)
+    encoder_obj = Encoder(
+        encoder=vocab,
+        bpe_merges=bpe_merges,
+        byte_encoder=byte_encoder,
+        byte_decoder=byte_decoder,
+        tokenize=tokenize
+    )
+    return encoder_obj
+def codegen_tokenize(text, tokenizer):
+    return tokenizer.encode(text)
+def codegen_decode(tokens, tokenizer):
+    return tokenizer.decode(tokens)