Hhhh

Running

App Files Files Community

Kfjjdjdjdhdhd commited on Mar 13

Commit

7b74407

verified ·

1 Parent(s): caa5001

Upload 26 files

Browse files

Files changed (22) hide show

api.py +422 -509
background_tasks.py +37 -110
codegen_api.py +8 -17
coder.py +139 -0
image_to_3d_api.py +13 -25
imagegen_api.py +7 -23
main.py +9 -56
model_loader.py +229 -725
models.py +53 -42
musicgen_api.py +9 -28
sadtalker_api.py +16 -183
sadtalker_utils.py +209 -820
sentiment_api.py +7 -20
stt_api.py +11 -27
summarization_api.py +8 -21
text_generation.py +93 -194
text_to_video_api.py +13 -25
tokenxxx.py +44 -114
translation_api.py +7 -16
tts_api.py +9 -20
xtts_api.py +21 -0
xxx.py +43 -114

api.py CHANGED Viewed

@@ -1,509 +1,422 @@
-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.functional as F
-import torchaudio
-import numpy as np
-from PIL import Image
-import io
-import tempfile
-import queue
-import json
-import base64
-from markupsafe import Markup
-from markupsafe import escape
-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 Conversational Avatar</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"/>
-    <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, .form-group input[type="text"] {
-            width: 100%;
-            padding: 10px;
-            border: 1px solid #ccc;
-            border-radius: 5px;
-            font-size: 16px;
-            box-sizing: border-box;
-        }
-        button, #recordButton, #stopButton {
-            padding: 10px 15px;
-            border: none;
-            border-radius: 5px;
-            background-color: #007bff;
-            color: white;
-            font-size: 18px;
-            cursor: pointer;
-            transition: background-color 0.3s ease;
-            margin-right: 5px;
-        }
-        button:hover, #recordButton:hover, #stopButton: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: 300px;
-        }
-        #videoOutput {
-            margin-top: 20px;
-            border: 1px solid #ddd;
-            border-radius: 5px;
-            overflow: hidden;
-        }
-        #videoOutput video {
-            width: 100%;
-            display: block;
-        }
-        #animatedText {
-            position: fixed;
-            top: 20px;
-            left: 20px;
-            font-size: 1.5em;
-            color: rgba(0, 0, 0, 0.1);
-            pointer-events: none;
-            z-index: -1;
-        }
-        #transcriptionOutput {
-            margin-top: 10px;
-            padding: 10px;
-            border: 1px solid #ddd;
-            border-radius: 5px;
-            background-color: #f9f9f9;
-            font-size: 14px;
-            word-break: break-word;
-        }
-        @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, #recordButton, #stopButton {
-                font-size: 16px;
-                padding: 8px 12px;
-            }
-            #output, #transcriptionOutput {
-                font-size: 14px;
-                padding: 10px;
-                margin-top: 15px;
-            }
-        }
-    </style>
-</head>
-<body>
-<div id="animatedText" class="animated-text animate__animated animate__fadeIn animate__infinite infinite">AI POWERED</div>
-<div class="container">
-    <div class="header animate__animated animate__fadeInDown">
-        <h1>Conversational Avatar</h1>
-    </div>
-    <div class="form-group animate__animated animate__fadeInLeft">
-        <textarea id="textInput" rows="3" placeholder="Or type your request here"></textarea>
-    </div>
-    <div class="form-group animate__animated animate__fadeInRight" style="text-align: center;">
-        <button onclick="generateResponse()" class="animate__animated animate__fadeInUp">Generate Avatar Response</button>
-    </div>
-    <div style="text-align: center; margin-bottom: 15px;">
-        <button id="recordButton" class="animate__animated animate__fadeInUp"><i class="fas fa-microphone"></i> Start Recording</button>
-        <button id="stopButton" class="animate__animated animate__fadeInUp" disabled><i class="fas fa-stop-circle"></i> Stop Recording</button>
-    </div>
-    <div id="transcriptionOutput" class="animate__animated animate__fadeIn">
-        <strong>Transcription:</strong>
-        <span id="transcriptionText"></span>
-    </div>
-    <div id="output" class="animate__animated animate__fadeIn">
-        <strong>Response:</strong><br>
-        <span id="responseText"></span>
-    </div>
-    <div id="videoOutput" class="animate__animated animate__fadeIn">
-        <video id="avatarVideo" controls></video>
-    </div>
-</div>
-<script>
-    let mediaRecorder;
-    let audioChunks = [];
-    let lastResponse = "";
-    let accumulatedText = "";
-    let eventSource = null;
-    let audioURL;
-    const recordButton = document.getElementById('recordButton');
-    const stopButton = document.getElementById('stopButton');
-    const transcriptionTextSpan = document.getElementById('transcriptionText');
-    const responseTextSpan = document.getElementById('responseText');
-    const avatarVideoPlayer = document.getElementById('avatarVideo');
-    const textInputField = document.getElementById('textInput');
-    recordButton.onclick = async () => {
-        try {
-            const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
-            mediaRecorder = new MediaRecorder(stream);
-            audioChunks = [];
-            mediaRecorder.ondataavailable = event => {
-                audioChunks.push(event.data);
-            };
-            mediaRecorder.onstop = async () => {
-                const audioBlob = new Blob(audioChunks, { type: 'audio/wav' });
-                const formData = new FormData();
-                formData.append('audio', audioBlob, 'recording.wav');
-                transcriptionTextSpan.innerText = "Transcribing...";
-                responseTextSpan.innerText = "";
-                avatarVideoPlayer.src = "";
-                try {
-                    const sttResponse = await fetch('/api/v1/stt', {
-                        method: 'POST',
-                        body: formData
-                    });
-                    if (!sttResponse.ok) {
-                        throw new Error(\`HTTP error! status: ${sttResponse.status}\`);
-                    }
-                    const sttData = await sttResponse.json();
-                    const transcribedText = sttData.text;
-                    transcriptionTextSpan.innerText = transcribedText || "Transcription failed.";
-                    if (transcribedText) {
-                        await generateAvatarVideoResponse(transcribedText);
-                    }
-                } catch (error) {
-                    console.error("STT or subsequent error:", error);
-                    transcriptionTextSpan.innerText = "Transcription error.";
-                    responseTextSpan.innerText = "Error processing audio.";
-                } finally {
-                    recordButton.disabled = false;
-                    stopButton.disabled = true;
-                }
-            };
-            recordButton.disabled = true;
-            stopButton.disabled = false;
-            transcriptionTextSpan.innerText = "Recording...";
-            mediaRecorder.start();
-        } catch (error) {
-            console.error("Error accessing microphone:", error);
-            transcriptionTextSpan.innerText = "Microphone access denied or error.";
-            recordButton.disabled = false;
-            stopButton.disabled = true;
-        }
-    };
-    stopButton.onclick = () => {
-        if (mediaRecorder && mediaRecorder.state === "recording") {
-            transcriptionTextSpan.innerText = "Processing...";
-            mediaRecorder.stop();
-            recordButton.disabled = true;
-            stopButton.disabled = true;
-        }
-    };
-    async function generateResponse() {
-        const inputText = textInputField.value;
-        if (!inputText.trim()) {
-            alert("Please enter text or record audio.");
-            return;
-        }
-        transcriptionTextSpan.innerText = inputText;
-        await generateAvatarVideoResponse(inputText);
-    }
-    async function generateAvatarVideoResponse(inputText) {
-        responseTextSpan.innerText = "Generating response...";
-        avatarVideoPlayer.src = "";
-        accumulatedText = "";
-        lastResponse = "";
-        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
-        };
-        if (eventSource) {
-            eventSource.close();
-        }
-        eventSource = new EventSource('/api/v1/generate_stream?' + new URLSearchParams(requestData).toString());
-        eventSource.onmessage = async 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;
-                }
-                responseTextSpan.innerHTML = marked.parse(accumulatedText);
-                try {
-                    const ttsResponse = await fetch('/api/v1/tts', {
-                        method: 'POST',
-                        headers: {
-                            'Content-Type': 'application/json'
-                        },
-                        body: JSON.stringify({ text: currentResponse })
-                    });
-                    if (!ttsResponse.ok) {
-                        throw new Error(\`TTS HTTP error! status: ${ttsResponse.status}\`);
-                    }
-                    const ttsBlob = await ttsResponse.blob();
-                    audioURL = URL.createObjectURL(ttsBlob);
-                    const sadTalkerResponse = await fetch('/api/v1/sadtalker', {
-                        method: 'POST',
-                        body: new URLSearchParams({
-                            'source_image': './examples/source_image/full_body_female.png',
-                            'driven_audio': audioURL,
-                            'preprocess': 'full',
-                            'still_mode': false,
-                            'use_enhancer': true
-                        })
-                    });
-                    if (!sadTalkerResponse.ok) {
-                        throw new Error(\`SadTalker HTTP error! status: ${sadTalkerResponse.status}\`);
-                    }
-                    const sadTalkerData = await sadTalkerResponse.json();
-                    const videoURL = sadTalkerData.video_url;
-                    avatarVideoPlayer.src = videoURL;
-                } catch (ttsError) {
-                    console.error("TTS or SadTalker error:", ttsError);
-                    responseTextSpan.innerHTML += "<br><br>Error generating audio or video avatar.";
-                }
-                return;
-            }
-            accumulatedText += event.data;
-            let partialText = accumulatedText.replace("<|endoftext|>", "").replace(/\s+(?=[.,，。])/g, '').trim();
-            responseTextSpan.innerHTML = marked.parse(partialText);
-        };
-        eventSource.onerror = function(error) {
-            console.error("SSE error", error);
-            eventSource.close();
-            responseTextSpan.innerText = "Error generating response stream.";
-        };
-        const outputDiv = document.getElementById("output");
-        outputDiv.classList.add("show");
-    }
-</script>
-</body>
-</html>
-"""
-feedback_queue = queue.Queue()
-@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))
-    response_queue = queue.Queue()
-    reasoning_queue.put({
-        'text_input': text,
-        'temperature': temp,
-        'top_k': top_k,
-        'top_p': top_p,
-        'repetition_penalty': reppenalty,
-        'response_queue': response_queue
-    })
-    @stream_with_context
-    def event_stream():
-        while True:
-            output = response_queue.get()
-            if "error" in output:
-                yield "data: <ERROR>\n\n"
-                break
-            text_chunk = output.get("text")
-            if text_chunk:
-                for word in text_chunk.split(' '):
-                    clean_word = word.strip()
-                    if clean_word:
-                        yield "data: " + clean_word + "\n\n"
-                yield "data: <END_STREAM>\n\n"
-                break
-    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))
-    response_queue = queue.Queue()
-    reasoning_queue.put({
-        'text_input': text,
-        'temperature': temp,
-        'top_k': top_k,
-        'top_p': top_p,
-        'repetition_penalty': reppenalty,
-        'response_queue': response_queue
-    })
-    output = response_queue.get()
-    if "error" in output:
-        return jsonify({"error": output["error"]}), 500
-    result_text = output.get("text", "").strip()
-    return jsonify({"response": result_text})
-@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():
-    return tts_route()
-@app.route("/api/v1/stt", methods=["POST"])
-def stt_api():
-    return stt_route()
-@app.route("/api/v1/sentiment", methods=["POST"])
-def sentiment_api():
-    return sentiment_route()
-@app.route("/api/v1/imagegen", methods=["POST"])
-def imagegen_api():
-    return imagegen_route()
-@app.route("/api/v1/musicgen", methods=["POST"])
-def musicgen_api():
-    return musicgen_route()
-@app.route("/api/v1/translation", methods=["POST"])
-def translation_api():
-    return translation_route()
-@app.route("/api/v1/codegen", methods=["POST"])
-def codegen_api():
-    return codegen_route()
-@app.route("/api/v1/text_to_video", methods=["POST"])
-def text_to_video_api():
-    return text_to_video_route()
-@app.route("/api/v1/summarization", methods=["POST"])
-def summarization_api():
-    return summarization_route()
-@app.route("/api/v1/image_to_3d", methods=["POST"])
-def image_to_3d_api():
-    return image_to_3d_route()
-@app.route("/api/v1/sadtalker", methods=["POST"])
-def sadtalker():
-    from sadtalker_api import router as sadtalker_router
-    return sadtalker_router.create_video()
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)

+from main import *
+from tts_api import tts_api as tts_module_api
+from stt_api import stt_api as stt_module_api
+from sentiment_api import sentiment_api as sentiment_module_api
+from imagegen_api import imagegen_api as imagegen_module_api
+from musicgen_api import musicgen_api as musicgen_module_api
+from translation_api import translation_api as translation_module_api
+from codegen_api import codegen_api as codegen_module_api
+from text_to_video_api import text_to_video_api as text_to_video_module_api
+from summarization_api import summarization_api as summarization_module_api
+from image_to_3d_api import image_to_3d_api as image_to_3d_module_api
+from xtts_api import xtts_api as xtts_module_api
+from flask import Flask, request, jsonify, Response, send_file, stream_with_context
+from flask_cors import CORS
+import io
+import queue
+import base64
+import gradio as gr
+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 = 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()
+@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))
+    response_queue = queue.Queue()
+    reasoning_queue.put({
+        'text_input': text,
+        'temperature': temp,
+        'top_k': top_k,
+        'top_p': top_p,
+        'repetition_penalty': reppenalty,
+        'response_queue': response_queue
+    })
+    @stream_with_context
+    def event_stream():
+        while True:
+            output = response_queue.get()
+            if "error" in output:
+                yield "data: <ERROR>\n\n"
+                break
+            text_chunk = output.get("text")
+            if text_chunk:
+                for word in text_chunk.split(' '):
+                    clean_word = word.strip()
+                    if clean_word:
+                        yield "data: " + clean_word + "\n\n"
+                yield "data: <END_STREAM>\n\n"
+                break
+    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))
+    response_queue = queue.Queue()
+    reasoning_queue.put({
+        'text_input': text,
+        'temperature': temp,
+        'top_k': top_k,
+        'top_p': top_p,
+        'repetition_penalty': reppenalty,
+        'response_queue': response_queue
+    })
+    output = response_queue.get()
+    if "error" in output:
+        return jsonify({"error": output["error"]}), 500
+    result_text = output.get("text", "").strip()
+    return jsonify({"response": result_text})
+@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():
+    return tts_module_api()
+@app.route("/api/v1/stt", methods=["POST"])
+def stt_api():
+    return stt_module_api()
+@app.route("/api/v1/sentiment", methods=["POST"])
+def sentiment_api():
+    return sentiment_module_api()
+@app.route("/api/v1/imagegen", methods=["POST"])
+def imagegen_api():
+    return imagegen_module_api()
+@app.route("/api/v1/musicgen", methods=["POST"])
+def musicgen_api():
+    return musicgen_module_api()
+@app.route("/api/v1/translation", methods=["POST"])
+def translation_api():
+    return translation_module_api()
+@app.route("/api/v1/codegen", methods=["POST"])
+def codegen_api():
+    return codegen_module_api()
+@app.route("/api/v1/text_to_video", methods=["POST"])
+def text_to_video_api():
+    return text_to_video_module_api()
+@app.route("/api/v1/summarization", methods=["POST"])
+def summarization_api():
+    return summarization_module_api()
+@app.route("/api/v1/image_to_3d", methods=["POST"])
+def image_to_3d_api():
+    return image_to_3d_module_api()
+@app.route("/api/v1/xtts_clone", methods=["POST"])
+def xtts_clone_api():
+    return xtts_module_api()
+@app.route("/api/v1/sadtalker", methods=["POST"])
+def sadtalker():
+    from sadtalker_api import router as sadtalker_router
+    return sadtalker_router.create_video()
+if __name__ == "__main__":
+    with gr.Blocks() as demo:
+        gr.Markdown("## AI Powerhouse")
+        with gr.Tab("Text Generation"):
+            text_input = gr.Textbox(lines=5, placeholder="Enter text")
+            text_output = gr.Markdown()
+            text_button = gr.Button("Generate Text")
+            text_button.click(generate, inputs=text_input, outputs=text_output)
+        with gr.Tab("Image Generation"):
+            image_text_input = gr.Textbox(lines=3, placeholder="Enter prompt for image")
+            image_output = gr.Image()
+            image_button = gr.Button("Generate Image")
+            image_button.click(imagegen_api, inputs=image_text_input, outputs=image_output)
+        with gr.Tab("Music Generation"):
+            music_text_input = gr.Textbox(lines=3, placeholder="Enter prompt for music")
+            music_output = gr.Audio()
+            music_button = gr.Button("Generate Music")
+            music_button.click(musicgen_api, inputs=music_text_input, outputs=music_output)
+        with gr.Tab("Code Generation"):
+            code_text_input = gr.Textbox(lines=3, placeholder="Enter prompt for code")
+            code_output = gr.File()
+            code_button = gr.Button("Generate Code")
+            code_button.click(codegen_api, inputs=code_text_input, outputs=code_output)
+        with gr.Tab("Text to Video"):
+            video_text_input = gr.Textbox(lines=3, placeholder="Enter prompt for video")
+            video_output = gr.Video()
+            video_button = gr.Button("Generate Video")
+            video_button.click(text_to_video_api, inputs=video_text_input, outputs=video_output)
+        with gr.Tab("Summarization"):
+            summary_text_input = gr.Textbox(lines=5, placeholder="Enter text to summarize")
+            summary_output = gr.Textbox()
+            summary_button = gr.Button("Summarize")
+            summary_button.click(summarization_api, inputs=summary_text_input, outputs=summary_output)
+        with gr.Tab("Translation"):
+            translate_text_input = gr.Textbox(lines=3, placeholder="Enter text to translate")
+            translate_lang_dropdown = gr.Dropdown(['es', 'en', 'fr', 'de'], value='es', label="Target Language")
+            translation_output = gr.Textbox()
+            translate_button = gr.Button("Translate")
+            translate_button.click(translation_api, inputs=[translate_text_input, translate_lang_dropdown], outputs=translation_output)
+        with gr.Tab("Sentiment Analysis"):
+            sentiment_text_input = gr.Textbox(lines=3, placeholder="Enter text for sentiment analysis")
+            sentiment_output = gr.Textbox()
+            sentiment_button = gr.Button("Analyze Sentiment")
+            sentiment_button.click(sentiment_api, inputs=sentiment_text_input, outputs=sentiment_output)
+        with gr.Tab("Text to Speech"):
+            tts_text_input = gr.Textbox(lines=3, placeholder="Enter text for speech")
+            tts_output = gr.Audio()
+            tts_button = gr.Button("Generate Speech")
+            tts_button.click(tts_api, inputs=tts_text_input, outputs=tts_output)
+        with gr.Tab("Voice Cloning (XTTS)"):
+            xtts_text_input = gr.Textbox(lines=3, placeholder="Enter text for voice cloning")
+            xtts_audio_input = gr.Audio(source="upload", type="filepath", label="Reference Audio for Voice Cloning")
+            xtts_output = gr.Audio()
+            xtts_button = gr.Button("Clone Voice")
+            xtts_button.click(xtts_module_api, inputs=[xtts_text_input, xtts_audio_input], outputs=xtts_output)
+        with gr.Tab("Speech to Text"):
+            stt_audio_input = gr.Audio(source="microphone", type="filepath")
+            stt_output = gr.Textbox()
+            stt_button = gr.Button("Transcribe Speech")
+            stt_button.click(stt_api, inputs=stt_audio_input, outputs=stt_output)
+        with gr.Tab("Image to 3D"):
+            image_3d_input = gr.Image(source="upload", type="filepath")
+            model_3d_output = gr.File()
+            image_3d_button = gr.Button("Generate 3D Model")
+            image_3d_button.click(image_to_3d_api, inputs=image_3d_input, outputs=model_3d_output)
+    app = gr.routes.App(demo)
+    app.run(host="0.0.0.0", port=7860)

background_tasks.py CHANGED Viewed

@@ -1,18 +1,9 @@
-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()
@@ -25,7 +16,7 @@ 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):
@@ -34,91 +25,46 @@ class SimpleClassifier(nn.Module):
         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).unsqueeze(0)
 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)
 def perform_reasoning_stream(text_input, temperature=0.7, top_k=40, top_p=0.0, repetition_penalty=1.2):
     for token in sample_sequence(text_input, model_gpt2, enc, length=999999999, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, device=device):
-        if token == "<END_STREAM>":
-            yield "<END_STREAM>"
-            break
         yield token + " "
 def background_reasoning_queue():
@@ -126,40 +72,21 @@ def background_reasoning_queue():
     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()

+import time, threading, queue, uuid, unicodedata, re
 from deep_translator import GoogleTranslator
 from duckduckgo_search import DDGS
+import nltk, torch, torch.nn as nn
 nltk.download('punkt')
 categories = ['News', 'Sports', 'Entertainment']
 TEXT_GENERATION_RATE = 10
 text_queue = queue.Queue()
 class SimpleClassifier(nn.Module):
     def __init__(self, vocab_size, num_classes, embedding_dim=128):
+        super().__init__()
         self.embedding = nn.Embedding(vocab_size, embedding_dim)
         self.fc = nn.Linear(embedding_dim, num_classes)
     def forward(self, x):
         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).unsqueeze(0)
 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: 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: time.sleep(5)
 def perform_reasoning_stream(text_input, temperature=0.7, top_k=40, top_p=0.0, repetition_penalty=1.2):
     for token in sample_sequence(text_input, model_gpt2, enc, length=999999999, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, device=device):
+        if token == "<END_STREAM>": yield "<END_STREAM>"; break
         yield token + " "
 def background_reasoning_queue():
     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: pass
+            if reasoning_queue and not reasoning_queue.empty(): reasoning_queue.task_done()

codegen_api.py CHANGED Viewed

@@ -1,22 +1,13 @@
 from flask import jsonify, send_file, request
 from main import *
 def generate_code(prompt, output_path="output_code.py"):
-    if codegen_model is None or codegen_tokenizer is None:
-        return "Code generation model or tokenizer not initialized."
-    input_ids = codegen_tokenizer(prompt, return_tensors='pt').to(device)
-    output = codegen_model.generate(input_ids, max_length=2048, 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 or tokenizer 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")

 from flask import jsonify, send_file, request
 from main import *
+import io, base64
 def generate_code(prompt, output_path="output_code.py"):
+    if codegen_model is None or codegen_tokenizer is None: return {"error": "Code generation model or tokenizer not initialized."}
+    input_ids = codegen_tokenizer(prompt, return_tensors='pt').to(device); output = codegen_model.generate(input_ids, max_length=2048, temperature=0.7, top_p=0.9)
+    code = codegen_tokenizer.decode(output[0], skip_special_tokens=True); return {"code": code}
+def codegen_api(prompt):
+    output = generate_code(prompt)
+    if "error" in output: return {"error": output["error"]}
+    code_base64 = base64.b64encode(output['code'].encode('utf-8')).decode('utf-8'); return {"code_base64": code_base64, "mimetype": "text/x-python", "filename": "output.py"}

coder.py ADDED Viewed

	@@ -0,0 +1,139 @@

+import dash
+from dash import Dash, html, dcc, callback, Output, Input, State
+from dash.exceptions import PreventUpdate
+import base64, uuid, requests, io, re
+app_style = None
+external_stylesheets = ['style.css']
+app = dash.Dash(__name__, external_stylesheets=external_stylesheets)
+app.layout = html.Div([dcc.Location(id='url', refresh=False), html.Div(id='page-content')])
+index_page = html.Div([html.H1("AI Powerhouse", className='index-title animate__animated animate__fadeInDown'), html.Div([dcc.Link('Text Generation', href='/text-generation', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Audio Video', href='/audio-video', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Image Generation', href='/image-generation', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Image to 3D', href='/image-to-3d', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Text to Video', href='/text-to-video', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Music Generation', href='/music-generation', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Sentiment Analysis', href='/sentiment-analysis', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Translation', href='/translation', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Code Generation', href='/code-generation', className='index-link animate__animated animate__fadeInUp'), dcc.Link('Summarization', href='/summarization', className='index-link animate__animated animate__fadeInUp'), dcc.Link('SadTalker', href='/sad-talker', className='index-link animate__animated animate__fadeInUp')], className='index-links-container')], className='index-page page-layout')
+text_generation_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Text Generation'), html.Div(className='chat-container page-layout', children=[html.Div(className='chat-header animate__animated animate__fadeInDown', children="Text Generation Interface"), html.Div(className='chat-form animate__animated animate__fadeInLeft', children=[dcc.Textarea(id='text-input', placeholder='Enter text prompt...', rows=4, className='chat-text-area'), html.Button('Generate', id='generate-button', n_clicks=0, className='chat-generate-button')]), html.Div(id='output', className='chat-output animate__animated animate__fadeInUp', children=[html.Div(className='response-header', children="Response:"), dcc.Markdown(id='generated-text', className='response-text')]), dcc.Link('Back to Home', href='/', className='chat-back-link')])], className='page-layout text-gen-layout')
+audio_video_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Audio & Video Tools'), html.Div(className='av-container page-layout', children=[html.Div(className='av-header animate__animated animate__fadeInDown', children="Audio & Video Processing"), html.Div(className='av-upload-section animate__animated animate__fadeInLeft', children=[html.Div(className='upload-box', children=[dcc.Upload(id='upload-audio', children=html.Div(['Drag and Drop or ', html.A('Select Audio')]), className='upload-area'), html.Div(id='audio-output-text', className='upload-output', children="STT Output")]), html.Div(className='upload-box', children=[dcc.Upload(id='upload-image', children=html.Div(['Drag and Drop or ', html.A('Select Image')]), className='upload-area'), html.Div(id='image-output', className='upload-output', children="Image Uploaded")])]), html.Div(id='video-output', className='av-video-output animate__animated animate__fadeInUp', children="Video Output Area"), dcc.Link('Back to Home', href='/', className='av-back-link')])], className='page-layout av-layout')
+image_generation_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Image Generation'), html.Div(className='imagegen-container page-layout', children=[html.Div(className='imagegen-header animate__animated animate__fadeInDown', children="Image Generation Interface"), html.Div(className='imagegen-form animate__animated animate__fadeInLeft', children=[dcc.Textarea(id='imagegen-text-input', placeholder='Enter prompt for image...', rows=4, className='imagegen-text-area'), html.Button('Generate Image', id='generate-image-button', n_clicks=0, className='imagegen-generate-button')]), html.Div(id='image-output-display', className='imagegen-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='imagegen-back-link')])], className='page-layout imagegen-layout')
+image_to_3d_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Image to 3D Conversion'), html.Div(className='imagetod-container page-layout', children=[html.Div(className='imagetod-header animate__animated animate__fadeInDown', children="Image to 3D Model Conversion"), html.Div(className='imagetod-upload animate__animated animate__fadeInLeft', children=[dcc.Upload(id='upload-image-3d', children=html.Div(['Drag and Drop or ', html.A('Select Image')]), className='imagetod-upload-area')]), html.Div(id='3d-model-output', className='imagetod-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='imagetod-back-link')])], className='page-layout imagetod-layout')
+text_to_video_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Text to Video Generation'), html.Div(className='textvideo-container page-layout', children=[html.Div(className='textvideo-header animate__animated animate__fadeInDown', children="Text to Video Generation Interface"), html.Div(className='textvideo-form animate__animated animate__fadeInLeft', children=[dcc.Textarea(id='text-video-input', placeholder='Enter prompt for video...', rows=4, className='textvideo-text-area'), html.Button('Generate Video', id='generate-video-button', n_clicks=0, className='textvideo-generate-button')]), html.Div(id='video-gen-output', className='textvideo-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='textvideo-back-link')])], className='page-layout textvideo-layout')
+music_generation_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Music Generation'), html.Div(className='musicgen-container page-layout', children=[html.Div(className='musicgen-header animate__animated animate__fadeInDown', children="Music Generation Interface"), html.Div(className='musicgen-form animate__animated animate__fadeInLeft', children=[dcc.Textarea(id='musicgen-text-input', placeholder='Enter prompt for music...', rows=4, className='musicgen-text-area'), html.Button('Generate Music', id='generate-music-button', n_clicks=0, className='musicgen-generate-button')]), html.Div(id='music-output-display', className='musicgen-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='musicgen-back-link')])], className='page-layout musicgen-layout')
+sentiment_analysis_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Sentiment Analysis'), html.Div(className='sentiment-container page-layout', children=[html.Div(className='sentiment-header animate__animated animate__fadeInDown', children="Sentiment Analysis Tool"), html.Div(className='sentiment-form animate__animated animate__fadeInLeft', children=[dcc.Textarea(id='sentiment-text-input', placeholder='Enter text for analysis...', rows=4, className='sentiment-text-area'), html.Button('Analyze Sentiment', id='analyze-sentiment-button', n_clicks=0, className='sentiment-analyze-button')]), html.Div(id='sentiment-output-display', className='sentiment-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='sentiment-back-link')])], className='page-layout sentiment-layout')
+translation_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Translation Services'), html.Div(className='translation-container page-layout', children=[html.Div(className='translation-header animate__animated animate__fadeInDown', children="Translation Interface"), html.Div(className='translation-form animate__animated animate__fadeInLeft', children=[dcc.Textarea(id='translate-text-input', placeholder='Enter text to translate...', rows=4, className='translation-text-area'), dcc.Dropdown(id='target-language-dropdown', options=[{'label': 'Spanish', 'value': 'es'},{'label': 'English', 'value': 'en'},{'label': 'French', 'value': 'fr'},{'label': 'German', 'value': 'de'}], value='es', className='translation-dropdown'), html.Button('Translate', id='translate-button', n_clicks=0, className='translation-translate-button')]), html.Div(id='translation-output-display', className='translation-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='translation-back-link')])], className='page-layout translation-layout')
+code_generation_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Code Generation'), html.Div(className='codegen-container page-layout', children=[html.Div(className='codegen-header animate__animated animate__fadeInDown', children="Code Generation Interface"), html.Div(className='codegen-form animate__animated animate__fadeInLeft', children=[dcc.Textarea(id='codegen-text-input', placeholder='Enter prompt for code...', rows=4, className='codegen-text-area'), html.Button('Generate Code', id='generate-code-button', n_clicks=0, className='codegen-generate-button')]), html.Div(id='codegen-output-display', className='codegen-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='codegen-back-link')])], className='page-layout codegen-layout')
+summarization_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED Text Summarization'), html.Div(className='summarization-container page-layout', children=[html.Div(className='summarization-header animate__animated animate__fadeInDown', children="Text Summarization Tool"), html.Div(className='summarization-form animate__animated animate__fadeInLeft', children=[dcc.Textarea(id='summarize-text-input', placeholder='Enter text to summarize...', rows=4, className='summarization-text-area'), html.Button('Summarize', id='summarize-button', n_clicks=0, className='summarization-summarize-button')]), html.Div(id='summarization-output-display', className='summarization-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='summarization-back-link')])], className='page-layout summarization-layout')
+sadtalker_layout = html.Div([html.Div(className='animated-text animate__animated animate__fadeIn animate__infinite infinite', children='AI POWERED SadTalker'), html.Div(className='sadtalker-container page-layout', children=[html.Div(className='sadtalker-header animate__animated animate__fadeInDown', children="SadTalker Interface"), html.Div(className='sadtalker-upload animate__animated animate__fadeInLeft', children=[dcc.Upload(id='upload-sadtalker-image', children=html.Div(['Drag and Drop Image', html.Br(), html.I(className="fas fa-image upload-icon")]), className='sadtalker-upload-area', multiple=False), dcc.Upload(id='upload-sadtalker-audio', children=html.Div(['Drag and Drop Audio', html.Br(), html.I(className="fas fa-file-audio upload-icon")]), className='sadtalker-upload-area', multiple=False)]), html.Div(id='sadtalker-video-output', className='sadtalker-output animate__animated animate__fadeInUp'), dcc.Link('Back to Home', href='/', className='sadtalker-back-link')])], className='page-layout sadtalker-layout')
+@app.callback(Output('page-content', 'children'), [Input('url', 'pathname')])
+def display_page(pathname):
+    if pathname == '/text-generation': return text_generation_layout
+    elif pathname == '/audio-video': return audio_video_layout
+    elif pathname == '/image-generation': return image_generation_layout
+    elif pathname == '/image-to-3d': return image_to_3d_layout
+    elif pathname == '/text-to-video': return text_to_video_layout
+    elif pathname == '/music-generation': return music_generation_layout
+    elif pathname == '/sentiment-analysis': return sentiment_analysis_layout
+    elif pathname == '/translation': return translation_layout
+    elif pathname == '/code-generation': return code_generation_layout
+    elif pathname == '/summarization': return summarization_layout
+    elif pathname == '/sad-talker': return sadtalker_layout
+    else: return index_page
+@app.callback(Output('generated-text', 'children'), Output('output', 'className'), Input('generate-button', 'n_clicks'), State('text-input', 'value'), prevent_initial_call=True)
+def generate_reasoning_dash(n_clicks, text_input):
+    if not text_input: return "Please enter text.", 'chat-output animate__animated animate__fadeInUp error-output'
+    api_url = "/api/v1/generate"; payload = {"text": text_input}
+    try: response = requests.post("http://127.0.0.1:7860" + api_url, json=payload); response.raise_for_status(); data = response.json(); generated_text = data.get("response", "Error in backend response."); return generated_text, 'chat-output animate__animated animate__fadeInUp'
+    except requests.exceptions.RequestException as e: return f"Error communicating with backend: {e}", 'chat-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('audio-output-text', 'children'), Output('video-output', 'children'), Output('audio-output-text', 'className'), Output('video-output', 'className'), Input('upload-audio', 'contents'), State('upload-audio', 'filename'), Input('upload-image', 'contents'), State('upload-image', 'filename'), prevent_initial_call=True)
+def process_audio_video_dash(audio_contents, audio_filename, image_contents, image_filename):
+    stt_output_text = ""; video_display = ""; stt_class = 'upload-output'; video_class = 'av-video-output animate__animated animate__fadeInUp'
+    if audio_contents:
+        try: content_type, content_string = audio_contents.split(','); decoded_audio = base64.b64decode(content_string); audio_io = io.BytesIO(decoded_audio); files = {'audio': (audio_filename, audio_io, content_type)}; response = requests.post("http://127.0.0.1:7860/api/v1/stt", files=files); response.raise_for_status(); data = response.json(); stt_output_text = f"STT Output: {data.get('text', 'Transcription failed')}"; stt_class = 'upload-output success'
+        except requests.exceptions.RequestException as e: stt_output_text = f"STT Error: {e}"; stt_class = 'upload-output error'
+    if image_contents:
+        try: content_type, content_string = image_contents.split(','); decoded_image = base64.b64decode(content_string); image_io = io.BytesIO(decoded_image); files = {'image': (image_filename, image_io, content_type)}; response = requests.post("http://127.0.0.1:7860/api/v1/image_to_3d", files=files); response.raise_for_status(); video_display = "3D Model Feature Extracted (Check Backend Logs for Output)."; video_class = 'av-video-output animate__animated animate__fadeInUp success'
+        except requests.exceptions.RequestException as e: video_display = f"3D Error: {e}"; video_class = 'av-video-output animate__animated animate__fadeInUp error'
+    video_output_component = html.Div(video_display) if video_display else ""; return stt_output_text, video_output_component, stt_class, video_class
+@app.callback(Output('image-output-display', 'children'), Output('image-output-display', 'className'), Input('generate-image-button', 'n_clicks'), State('imagegen-text-input', 'value'), prevent_initial_call=True)
+def generate_image_dash(n_clicks, prompt):
+    if not prompt: return "Please enter a prompt for image generation.", 'imagegen-output animate__animated animate__fadeInUp error-output'
+    api_url = "/api/v1/imagegen"; payload = {"prompt": prompt}
+    try: response = requests.post("http://127.0.0.1:7860" + api_url, json=payload); response.raise_for_status(); image_base64 = base64.b64encode(response.content).decode('utf-8'); return html.Img(src=f'data:image/png;base64,{image_base64}', className='generated-image'), 'imagegen-output animate__animated animate__fadeInUp success-output'
+    except requests.exceptions.RequestException as e: return f"Image Generation Error: {e}", 'imagegen-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('3d-model-output', 'children'), Output('3d-model-output', 'className'), Input('upload-image-3d', 'contents'), State('upload-image-3d', 'filename'), prevent_initial_call=True)
+def process_image_to_3d_dash(contents, filename):
+    if contents is None: raise PreventUpdate
+    try:
+        content_type, content_string = contents.split(',')
+        decoded_image = base64.b64decode(content_string); image_io = io.BytesIO(decoded_image); files = {'image': (filename, image_io, content_type)}
+        response = requests.post("http://127.0.0.1:7860/api/v1/image_to_3d", files=files); response.raise_for_status()
+        content_disposition = response.headers.get('Content-Disposition'); download_filename = 'model_3d.obj'
+        if content_disposition: filenames = re.findall('filename="([^"]+)"', content_disposition);
+        if filenames: download_filename = filenames[0]
+        model_base64 = base64.b64encode(response.content).decode('utf-8'); href = f'data:model/obj;base64,{model_base64}'
+        download_link = html.A('Download 3D Model', href=href, download=download_filename, className='download-link'),; return download_link, 'imagetod-output animate__animated animate__fadeInUp success-output'
+    except requests.exceptions.RequestException as e: return f"3D Conversion Error: {e}", 'imagetod-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('video-gen-output', 'children'), Output('video-gen-output', 'className'), Input('generate-video-button', 'n_clicks'), State('text-video-input', 'value'), prevent_initial_call=True)
+def generate_video_dash(n_clicks, prompt):
+    if not prompt: return "Please enter a prompt for video generation.", 'textvideo-output animate__animated animate__fadeInUp error-output'
+    api_url = "/api/v1/text_to_video"; payload = {"prompt": prompt}
+    try: response = requests.post("http://127.0.0.1:7860" + api_url, json=payload); response.raise_for_status(); video_base64 = base64.b64encode(response.content).decode('utf-8'); return html.Video(src=f'data:video/mp4;base64,{video_base64}', controls=True, className='generated-video'), 'textvideo-output animate__animated animate__fadeInUp success-output'
+    except requests.exceptions.RequestException as e: return f"Video Generation Error: {e}", 'textvideo-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('music-output-display', 'children'), Output('music-output-display', 'className'), Input('generate-music-button', 'n_clicks'), State('musicgen-text-input', 'value'), prevent_initial_call=True)
+def generate_music_dash(n_clicks, prompt):
+    if not prompt: return "Please enter a prompt for music generation.", 'musicgen-output animate__animated animate__fadeInUp error-output'
+    api_url = "/api/v1/musicgen"; payload = {"prompt": prompt}
+    try: response = requests.post("http://127.0.0.1:7860" + api_url, json=payload); response.raise_for_status(); audio_base64 = base64.b64encode(response.content).decode('utf-8'); return html.Audio(src=f'data:audio/wav;base64,{audio_base64}', controls=True, className='generated-audio'), 'musicgen-output animate__animated animate__fadeInUp success-output'
+    except requests.exceptions.RequestException as e: return f"Music Generation Error: {e}", 'musicgen-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('sentiment-output-display', 'children'), Output('sentiment-output-display', 'className'), Input('analyze-sentiment-button', 'n_clicks'), State('sentiment-text-input', 'value'), prevent_initial_call=True)
+def analyze_sentiment_dash(n_clicks, text):
+    if not text: return "Please enter text for sentiment analysis.", 'sentiment-output animate__animated animate__fadeInUp error-output'
+    api_url = "/api/v1/sentiment"; payload = {"text": text}
+    try: response = requests.post("http://127.0.0.1:7860" + api_url, json=payload); response.raise_for_status(); data = response.json(); sentiment_label = data.get('sentiment', 'Analysis Failed'); return f"Sentiment: {sentiment_label}", 'sentiment-output animate__animated animate__fadeInUp success-output'
+    except requests.exceptions.RequestException as e: return f"Sentiment Analysis Error: {e}", 'sentiment-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('translation-output-display', 'children'), Output('translation-output-display', 'className'), Input('translate-button', 'n_clicks'), State('translate-text-input', 'value'), State('target-language-dropdown', 'value'), prevent_initial_call=True)
+def translate_text_dash(n_clicks, text, target_lang):
+    if not text: return "Please enter text for translation.", 'translation-output animate__animated animate__fadeInUp error-output'
+    api_url = "/api/v1/translation"; payload = {"text": text, "target_lang": target_lang}
+    try: response = requests.post("http://127.0.0.1:7860" + api_url, json=payload); response.raise_for_status(); data = response.json(); translation = data.get('translated_text', 'Translation Failed'); return f"Translation ({target_lang.upper()}): {translation}", 'translation-output animate__animated animate__fadeInUp success-output'
+    except requests.exceptions.RequestException as e: return f"Translation Error: {e}", 'translation-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('codegen-output-display', 'children'), Output('codegen-output-display', 'className'), Input('generate-code-button', 'n_clicks'), State('codegen-text-input', 'value'), prevent_initial_call=True)
+def generate_code_dash(n_clicks, prompt):
+    if not prompt: return "Please enter a prompt for code generation.", 'codegen-output animate__animated animate__fadeInUp error-output'
+    api_url = "/api/v1/codegen"; payload = {"prompt": prompt}
+    try: response = requests.post("http://127.0.0.1:7860" + api_url, json=payload); response.raise_for_status()
+        content_disposition = response.headers.get('Content-Disposition'); download_filename = 'code.py'
+        if content_disposition: filenames = re.findall('filename="([^"]+)"', content_disposition); if filenames: download_filename = filenames[0]
+        code_base64 = base64.b64encode(response.content).decode('utf-8'); download_link = html.A('Download Code', href=f'data:text/x-python;base64,{code_base64}', download=download_filename, className='download-link'); return download_link, 'codegen-output animate__animated animate__fadeInUp success-output'
+    except requests.exceptions.RequestException as e: return f"Code Generation Error: {e}", 'codegen-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('summarization-output-display', 'children'), Output('summarization-output-display', 'className'), Input('summarize-button', 'n_clicks'), State('summarize-text-input', 'value'), prevent_initial_call=True)
+def summarize_text_dash(n_clicks, text):
+    if not text: return "Please enter text for summarization.", 'summarization-output animate__animated animate__fadeInUp error-output'
+    api_url = "/api/v1/summarization"; payload = {"text": text}
+    try: response = requests.post("http://127.0.0.1:7860" + api_url, json=payload); response.raise_for_status()
+        content_disposition = response.headers.get('Content-Disposition'); download_filename = 'summary.txt'
+        if content_disposition: filenames = re.findall('filename="([^"]+)"', content_disposition); if filenames: download_filename = filenames[0]
+        summary_base64 = base64.b64encode(response.content).decode('utf-8'); download_link = html.A('Download Summary', href=f'data:text/plain;base64,{summary_base64}', download=download_filename, className='download-link'); return download_link, 'summarization-output animate__animated animate__fadeInUp success-output'
+    except requests.exceptions.RequestException as e: return f"Summarization Error: {e}", 'summarization-output animate__animated animate__fadeInUp error-output'
+@app.callback(Output('sadtalker-video-output', 'children'), Output('sadtalker-video-output', 'className'), Input('upload-sadtalker-image', 'contents'), State('upload-sadtalker-image', 'filename'), Input('upload-sadtalker-audio', 'contents'), State('upload-sadtalker-audio', 'filename'), prevent_initial_call=True)
+def process_sadtalker_dash(image_contents, image_filename, audio_contents, audio_filename):
+    if not image_contents or not audio_contents: return "Please upload both image and audio for SadTalker.", 'sadtalker-output animate__animated animate__fadeInUp error-output'
+    try:
+        image_content_type, image_content_string = image_contents.split(','); decoded_image = base64.b64decode(image_content_string); image_io = io.BytesIO(decoded_image)
+        audio_content_type, audio_content_string = audio_contents.split(','); decoded_audio = base64.b64decode(audio_content_string); audio_io = io.BytesIO(decoded_audio)
+        files = {'source_image_file': (image_filename, image_io, image_content_type), 'driven_audio_file': (audio_filename, audio_io, audio_content_type)}
+        response = requests.post("http://127.0.0.1:7860/api/v1/sadtalker", files=files); response.raise_for_status(); data = response.json(); video_url = data.get('video_url')
+        if video_url: video_base64 = base64.b64encode(requests.get(video_url).content).decode('utf-8'); return html.Video(src=f'data:video/mp4;base64,{video_base64}', controls=True, className='generated-video'), 'sadtalker-output animate__animated animate__fadeInUp success-output'
+        else: return "SadTalker video generation failed, check backend logs.", 'sadtalker-output animate__animated animate__fadeInUp error-output'
+    except requests.exceptions.RequestException as e: return f"SadTalker Error: {e}", 'sadtalker-output animate__animated animate__fadeInUp error-output'
+if __name__ == '__main__': app.run_server(host='0.0.0.0', port=7861, debug=False)

image_to_3d_api.py CHANGED Viewed

@@ -1,31 +1,19 @@
-import os
-import uuid
 from flask import jsonify, send_file, request
 from main import *
 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")

+```
+```python
+--- START OF FILE image_to_3d_api.py ---
+import os, uuid
 from flask import jsonify, send_file, request
 from main import *
 from PIL import Image
+import torch, numpy as np, io, base64
 def image_to_3d_func(image_path, output_path="output_3d.obj"):
+    if image_to_3d_model is None: return {"error": "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); return {"model_3d": mesh_obj}
+def image_to_3d_api(image_path):
+    output = image_to_3d_func(image_path)
+    if "error" in output: return {"error": output["error"]}
+    model_3d_base64 = base64.b64encode(output['model_3d'].encode('utf-8')).decode('utf-8'); return {"model_3d_base64": model_3d_base64, "mimetype": "model/obj", "filename": "output_3d.obj"}

imagegen_api.py CHANGED Viewed

@@ -3,31 +3,15 @@ from flask import jsonify, send_file, request
 from io import BytesIO
 from PIL import Image
 from main import *
-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)
-    with torch.no_grad():
-        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")

 from io import BytesIO
 from PIL import Image
 from main import *
+import torch, base64
 def generate_image(prompt, output_path="output_image.png"):
+    if imagegen_model is None: return {"error": "Image generation model not initialized."}
     generator = torch.Generator(device=device).manual_seed(0)
+    with torch.no_grad(): image = imagegen_model(prompt, generator=generator,).images[0]; image.save(output_path); return output_path
+def imagegen_api(prompt):
     output_file = generate_image(prompt)
+    if isinstance(output_file, dict) and "error" in output_file: return {"error": output_file["error"]}
+    image_io = BytesIO(); pil_image = Image.open(output_file); pil_image.save(image_io, 'PNG'); image_base64 = base64.b64encode(image_io.getvalue()).decode('utf-8')
+    os.remove(output_file); return {"image_base64": image_base64, "mimetype": "image/png"}

main.py CHANGED Viewed

@@ -1,10 +1,4 @@
-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 *
@@ -19,50 +13,13 @@ from background_tasks import *
 from text_generation import *
 from sadtalker_utils import *
-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
     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)
@@ -71,6 +28,7 @@ def load_models():
     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)
     sadtalker_instance = SadTalker(checkpoint_path='./checkpoints', config_path='./src/config')
 if __name__ == "__main__":
@@ -78,13 +36,8 @@ if __name__ == "__main__":
     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)

+import threading, queue, time, os, nltk, re, json
 from flask import Flask
 from flask_cors import CORS
 from api import *
 from text_generation import *
 from sadtalker_utils import *
+state_dict, enc, config, model_gpt2, device, news_clf, tfidf_vectorizer, text_queue, categories, background_threads, feedback_queue, reasoning_queue, seen_responses, dialogue_history, vocabulary, word_to_index, index_to_word, translation_model, sp, codegen_model, codegen_tokenizer, codegen_vocabulary, codegen_index_to_word, codegen_word_to_index, summarization_model, summarization_vocabulary, summarization_word_to_index, summarization_index_to_word, sadtalker_instance, imagegen_model, image_to_3d_model, text_to_video_model, stream_type, sentiment_model, stt_model, tts_model, musicgen_model, xtts_model = None, None, None, None, torch.device("cuda" if torch.cuda.is_available() else "cpu"), None, None, queue.Queue(), None, [], queue.Queue(), queue.Queue(), set(), [], set(), {}, [], None, None, None, None, None, None, set(), {}, [], None, None, None, None, "text", None, None, None, None, None
 def load_models():
+    global model_gpt2, enc, translation_model, codegen_model, codegen_tokenizer, summarization_model, imagegen_model, image_to_3d_model, text_to_video_model, sadtalker_instance, sentiment_model, stt_model, tts_model, musicgen_model, xtts_model
     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, _, _, _ = 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)
     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)
+    xtts_model = initialize_xtts_model(XTTS_FOLDER, XTTS_FILES_URLS)
     sadtalker_instance = SadTalker(checkpoint_path='./checkpoints', config_path='./src/config')
 if __name__ == "__main__":
     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 CHANGED Viewed

@@ -1,725 +1,229 @@
-from tokenxxx import *
-from constants import *
-from utils import *
-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)
-    while not os.path.exists(filepath):
-        try:
-            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))
-        except Exception:
-            continue
-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()
-    merge_ranks = {}
-    for i, merge in enumerate(merges):
-        parts = merge.strip().split()
-        if len(parts) == 2:
-            merge_ranks[tuple(parts)] = i
-    def bpe(token):
-        word = list(token)
-        pairs = [(word[i], word[i+1]) for i in range(len(word)-1)]
-        while True:
-            candidate = None
-            candidate_rank = None
-            candidate_index = None
-            for i, pair in enumerate(pairs):
-                if pair in merge_ranks:
-                    rank = merge_ranks[pair]
-                    if candidate is None or rank < candidate_rank:
-                        candidate = pair
-                        candidate_rank = rank
-                        candidate_index = i
-            if candidate is None:
-                break
-            first, second = candidate
-            new_word = []
-            i = 0
-            while i < len(word):
-                if i < len(word) - 1 and word[i] == first and word[i+1] == second:
-                    new_word.append(first + second)
-                    i += 2
-                else:
-                    new_word.append(word[i])
-                    i += 1
-            word = new_word
-            if len(word) == 1:
-                break
-            pairs = [(word[i], word[i+1]) for i in range(len(word)-1)]
-        return word
-    def tokenizer(text):
-        tokens = []
-        for token in text.split():
-            bpe_tokens = bpe(token)
-            for subtoken in bpe_tokens:
-                tokens.append(vocab.get(subtoken, 0))
-        return tokens
-    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(x))
-        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)
-        self.alphas = 1 - self.betas
-        self.alpha_bars = torch.cumprod(self.alphas, dim=0)
-    def step(self, noise, t, sample):
-        alpha_bar = self.alpha_bars[t]
-        alpha_bar_prev = self.alpha_bars[t-1] if t > 0 else torch.tensor(1.0, device=sample.device)
-        x0 = (sample - torch.sqrt(1 - alpha_bar) * noise) / torch.sqrt(alpha_bar)
-        new_sample = torch.sqrt(alpha_bar_prev) * x0 + torch.sqrt(1 - alpha_bar_prev) * noise
-        return new_sample
-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

+from tokenxxx import *
+from constants import *
+from utils import *
+import os, json, urllib.request, urllib.parse, torch, hashlib, inspect
+from tqdm import tqdm
+from TTS.config import load_config
+from TTS.tts.models.xtts import Xtts
+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)
+    while not os.path.exists(filepath):
+        try:
+            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))
+        except: continue
+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()
+    merge_ranks = {};
+    for i, merge in enumerate(merges): parts = merge.strip().split(); if len(parts) == 2: merge_ranks[tuple(parts)] = i
+    def bpe(token):
+        word = list(token); pairs = [(word[i], word[i+1]) for i in range(len(word)-1)]
+        while True: candidate = None; candidate_rank = None; candidate_index = None
+            for i, pair in enumerate(pairs): if pair in merge_ranks: rank = merge_ranks[pair]; if candidate is None or rank < candidate_rank: candidate = pair; candidate_rank = rank; candidate_index = i
+            if candidate is None: break
+            first, second = candidate; 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
+            word = new_word;
+            if len(word) == 1: break
+            pairs = [(word[i], word[i+1]) for i in range(len(word)-1)]
+        return word
+    def tokenizer(text): tokens = []; for token in text.split(): bpe_tokens = bpe(token); for subtoken in bpe_tokens: tokens.append(vocab.get(subtoken, 0)); return tokens
+    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 XTTSConfig:
+    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(x)); 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__(); 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)
+        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); self.alphas = 1 - self.betas; self.alpha_bars = torch.cumprod(self.alphas, dim=0)
+    def step(self, noise, t, sample): alpha_bar = self.alpha_bars[t]; alpha_bar_prev = self.alpha_bars[t-1] if t > 0 else torch.tensor(1.0, device=sample.device)
+        x0 = (sample - torch.sqrt(1 - alpha_bar) * noise) / torch.sqrt(alpha_bar); new_sample = torch.sqrt(alpha_bar_prev) * x0 + torch.sqrt(1 - alpha_bar_prev) * noise; return new_sample
+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)
+class XTTSModelClass(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.xtts = XTTSModel(config, num_speakers=1024, num_languages=25)  # Adjust num_speakers, num_languages as needed
+    def forward(self, text_tokens, text_lengths, speaker_ids, language_ids, voice_samples, voice_sample_lengths):
+        return self.xtts.forward(text_tokens, text_lengths, speaker_ids, language_ids, voice_samples, voice_sample_lengths)
+    def inference(self, text, language_id, speaker_id, voice_sample, temperature=0.7, length_penalty=1.0):
+        return self.xtts.inference(text, language_id, speaker_id, voice_sample, temperature, length_penalty)
+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
+def initialize_xtts_model(folder, files):
+    download_files(folder, files)
+    config_xtts = XTTSConfig.from_dict(read_json(os.path.join(folder, "config.json")))
+    model = XTTSModelClass(config_xtts).to(device)
+    checkpoint = torch.load(os.path.join(folder, "model.pth"), map_location=torch.device(device))
+    model.load_state_dict(checkpoint["model"], strict=False)
+    model.eval()
+    return model.xtts

models.py CHANGED Viewed

@@ -2,9 +2,11 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
-import copy
-from configs import *
-from extensions import *
 class SentimentClassifierModel(nn.Module):
     def __init__(self, config):
@@ -13,15 +15,12 @@ class SentimentClassifierModel(nn.Module):
         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):
@@ -35,17 +34,11 @@ class STTModel(nn.Module):
         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):
@@ -55,15 +48,9 @@ class TTSModel(nn.Module):
         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):
@@ -72,23 +59,47 @@ class MusicGenModel(nn.Module):
         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

 import torch.nn as nn
 import torch.nn.functional as F
 import math
+from configs import MusicGenConfig
+from extensions import CodeGenBlock
+from TTS.tts.layers.xtts.transformer import XTransformerEncoder, XTransformerDecoder
+from TTS.tts.layers.xtts.flow import VitsFlowModules
+from TTS.tts.layers.xtts.tokenizer import VoiceBPE
 class SentimentClassifierModel(nn.Module):
     def __init__(self, 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):
         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(embedded, 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):
         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):
         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
+class XTTSModelClass(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.xtts = XTTSModel(config, num_speakers=1024, num_languages=25)
+    def forward(self, text_tokens, text_lengths, speaker_ids, language_ids, voice_samples, voice_sample_lengths):
+        return self.xtts.forward(text_tokens, text_lengths, speaker_ids, language_ids, voice_samples, voice_sample_lengths)
+    def inference(self, text, language_id, speaker_id, voice_sample, temperature=0.7, length_penalty=1.0):
+        return self.xtts.inference(text, language_id, speaker_id, voice_sample, temperature, length_penalty)
+class XTTSModel(nn.Module):
+    def __init__(self, config, num_speakers, num_languages):
+        super().__init__()
+        self.config = config
+        self.num_speakers = num_speakers
+        self.num_languages = num_languages
+        self.encoder = XTransformerEncoder(**config.encoder_config)
+        self.decoder = XTransformerDecoder(**config.decoder_config)
+        self.flow_modules = VitsFlowModules(**config.flow_config)
+        self.voice_tokenizer = VoiceBPE(vocab_path=config.voice_tokenizer_config.vocab_path, vocab_size=config.voice_tokenizer_config.vocab_size)
+        self.language_embedding = nn.Embedding(num_languages, config.embedding_dim)
+        self.speaker_embedding = nn.Embedding(num_speakers, config.embedding_dim)
+        self.text_embedding = nn.Embedding(config.num_chars, config.embedding_dim)
+    def forward(self, text_tokens, text_lengths, speaker_ids, language_ids, voice_samples, voice_sample_lengths):
+        lang_embed = self.language_embedding(language_ids); spk_embed = self.speaker_embedding(speaker_ids); text_embed = self.text_embedding(text_tokens)
+        encoder_outputs, _ = self.encoder(text_embed, text_lengths, lang_embed + spk_embed); mel_outputs, _ = self.decoder(encoder_outputs, lang_embed + spk_embed, voice_samples); return mel_outputs, None
+    def inference(self, text, language_id, speaker_id, voice_sample, temperature=0.7, length_penalty=1.0):
+        language_ids = torch.tensor([language_id], dtype=torch.long).to(device); speaker_ids = torch.tensor([speaker_id], dtype=torch.long).to(device)
+        text_tokens = self.voice_tokenizer.text_to_ids(text).to(device); text_lengths = torch.tensor([text_tokens.shape[0]], dtype=torch.long).to(device); voice_sample_lengths = torch.tensor([voice_sample.shape[0]], dtype=torch.long).to(device)
+        lang_embed = self.language_embedding(language_ids); spk_embed = self.speaker_embedding(speaker_ids); text_embed = self.text_embedding(text_tokens)
+        encoder_outputs, _ = self.encoder(text_embed, text_lengths, lang_embed + spk_embed); mel_outputs, _ = self.decoder.inference(encoder_outputs, lang_embed + spk_embed, voice_sample, temperature=temperature, length_penalty=length_penalty)
+        return mel_outputs

musicgen_api.py CHANGED Viewed

@@ -1,34 +1,15 @@
 from flask import jsonify, send_file, request
 from main import *
-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 = 10
-    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")

 from flask import jsonify, send_file, request
 from main import *
+import torch, soundfile as sf, numpy as np, io, base64
 def generate_music(prompt, output_path="output_music.wav"):
+    if musicgen_model is None: return {"error": "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(prompt):
     output_file = generate_music(prompt)
+    if isinstance(output_file, dict) and "error" in output_file: return {"error": output_file["error"]}
+    with open(output_file, 'rb') as f: audio_content = f.read()
+    audio_base64 = base64.b64encode(audio_content).decode('utf-8'); os.remove(output_file); return {"audio_base64": audio_base64, "mimetype": "audio/wav"}

sadtalker_api.py CHANGED Viewed

@@ -1,9 +1,4 @@
-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
@@ -19,186 +14,24 @@ 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(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)})
 router = APIRouter()
 router.add_api_route("/sadtalker", create_video, methods=["POST"])
-router.add_api_websocket_route("/ws", websocket_endpoint)

+import os, tempfile, uuid, asyncio, shutil, requests
 from urllib.parse import urlparse
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form, WebSocket
 from fastapi.responses import JSONResponse
 router = APIRouter()
 @router.post("/sadtalker")
+async def create_video(source_image_file: UploadFile = File(...), driven_audio_file: UploadFile = File(...)):
+    if not source_image_file: raise HTTPException(status_code=400, detail="Source image file is required")
+    if not driven_audio_file: raise HTTPException(status_code=400, detail="Driven audio file is required")
+    temp_source_image = tempfile.NamedTemporaryFile(suffix=os.path.splitext(source_image_file.filename)[1], delete=False)
+    content_image = await source_image_file.read(); temp_source_image.write(content_image); source_image_path = temp_source_image.name
+    temp_driven_audio = tempfile.NamedTemporaryFile(suffix=os.path.splitext(driven_audio_file.filename)[1], delete=False)
+    content_audio = await driven_audio_file.read(); temp_driven_audio.write(content_audio); driven_audio_path = temp_driven_audio.name
     try:
+        loop = asyncio.get_running_loop()
+        output_path = await loop.run_in_executor(None, sadtalker_instance.test, source_image_path, driven_audio_path)
+        video_base64 = None
+        with open(output_path, 'rb') as video_file: video_bytes = video_file.read(); video_base64 = base64.b64encode(video_bytes).decode('utf-8')
+        os.remove(output_path); return {"video_base64": video_base64, "mimetype": "video/mp4"}
+    except Exception as e: raise HTTPException(status_code=500, detail=str(e))
+    finally: os.remove(temp_source_image.name); os.remove(temp_driven_audio.name)
 router = APIRouter()
 router.add_api_route("/sadtalker", create_video, methods=["POST"])

sadtalker_utils.py CHANGED Viewed

@@ -1,820 +1,209 @@
-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
-import tqdm
-import math
-import torchaudio
-import urllib.request
-from safetensors.torch import load_file, save_file
-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))
-def load_state_dict_robust(model, checkpoint_path, device, model_name="model"):
-    if not os.path.exists(checkpoint_path):
-        raise FileNotFoundError(f"Checkpoint file not found: {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)
-    state_dict = checkpoint.get(model_name, checkpoint)
-    try:
-        model.load_state_dict(state_dict)
-    except RuntimeError as e:
-        print(f"Error loading {model_name} state_dict: {e}")
-        print(f"Trying to load state_dict with key mapping for {model_name}.")
-        model_state_dict = model.state_dict()
-        mapped_state_dict = {}
-        for key, value in state_dict.items():
-            if key in model_state_dict and model_state_dict[key].shape == value.shape:
-                mapped_state_dict[key] = value
-            else:
-                print(f"Skipping key {key} due to shape mismatch or missing in model.")
-        missing_keys, unexpected_keys = model.load_state_dict(mapped_state_dict, strict=False)
-        if missing_keys or unexpected_keys:
-            print(f"Missing keys: {missing_keys}")
-            print(f"Unexpected keys: {unexpected_keys}")
-        print(f"Successfully loaded {model_name} state_dict with key mapping.")
-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 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
-        for filename, url in [
-            (kp_file, kp_url), (aud_file, aud_url), (wav_file, wav_url), (gen_file, gen_url),
-            (mapx_file, mapx_url), (den_file, den_url), ('GFPGANv1.4.pth', GFPGAN_URL),
-            ('RealESRGAN_x2plus.pth', REALESRGAN_URL)
-        ]:
-            download_model(url, filename, 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)
-        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
-        face_enhancer = self.sadtalker_model.face_enhancer if self.use_enhancer else None
-        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
-            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
-        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')
-        load_state_dict_robust(self.kp_extractor, checkpoint_path, device, model_name='kp_detector')
-    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')
-        load_state_dict_robust(self.audio_model, checkpoint_path, device, model_name='wav2vec2')
-        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'], 'auido2pose_00140-model.pth')
-        load_state_dict_robust(self, mapping_checkpoint, device)
-    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']
-        face_3d = mapping(expression_coeff, pose_coeff, blink_coeff) if blink_coeff is not None else mapping(expression_coeff, pose_coeff)
-        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})
-        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']
-        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.forward(np.array(pil_image))
-                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')
-        load_state_dict_robust(self.generator, checkpoint_path, device, model_name='generator')
-    def forward(self, source_image, dense_motion, bg_param, face_3d_param=None):
-        video_3d = self.generator(source_image, kp_driving=dense_motion, bg_param=bg_param, face_3d_param=face_3d_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')
-        load_state_dict_robust(self.mapping_net, checkpoint_path, device, model_name='mapping')
-        self.f_3d_mean = torch.zeros(1, 64, device=device)
-    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')
-        load_state_dict_robust(self.dense_motion_network, checkpoint_path, device, model_name='dense_motion')
-    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):
-        if self.face_enhancer:
-            return self.face_enhancer.enhance(x, outscale=1)[0]
-        return x
-def load_models():
-    checkpoint_path = './checkpoints'
-    config_path = './src/config'
-    size = 256
-    preprocess = 'crop'
-    old_version = False
-    sadtalker_instance = SadTalker(checkpoint_path, config_path, size, preprocess, old_version)
-    print("SadTalker models loaded successfully!")
-    return sadtalker_instance
-if __name__ == '__main__':
-    sadtalker_instance = load_models()

+import os, shutil, uuid, cv2, numpy as np, torch, torch.nn as nn, torch.nn.functional as F, yaml, safetensors, librosa, imageio
+from PIL import Image
+from skimage import img_as_ubyte, transform
+from scipy.io import loadmat, wavfile
+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
+        for filename, url in [(kp_file, kp_url), (aud_file, aud_url), (wav_file, wav_url), (gen_file, gen_url), (mapx_file, mapx_url), (den_file, den_url), ('GFPGANv1.4.pth', GFPGAN_URL), ('RealESRGAN_x2plus.pth', REALESRGAN_URL)]: download_model(url, filename, checkpoint_dir)
+        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)
+        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, generator, mapping, he_estimator, audio_to_coeff, animate_from_coeff, face_enhancer = self.sadtalker_model.kp_extractor, self.sadtalker_model.generator, self.sadtalker_model.mapping, self.sadtalker_model.he_estimator, self.sadtalker_model.audio_to_coeff, self.sadtalker_model.animate_from_coeff, self.sadtalker_model.face_enhancer
+        with torch.no_grad():
+            kp_source = kp_extractor(batch['source_image'])
+            if self.still_mode or self.use_idle_mode: pose_coeff = audio_to_coeff.get_pose_coeff(batch['audio'], batch['ref_pose_coeff']); expression_coeff = audio_to_coeff.get_exp_coeff(batch['audio'], batch['ref_expression_coeff'])
+            elif self.use_idle_mode: pose_coeff = audio_to_coeff.get_pose_coeff(batch['audio'], batch['ref_pose_coeff']); expression_coeff = audio_to_coeff.get_exp_coeff(batch['audio'], batch['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
+            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.sadtalker = SadTalkerInnerModel(sadtalker_cfg, device_id)
+        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
+        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): return 0
+    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): return 0
+    def generate_idles_expression(self, length_of_audio): return 0
+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'); load_state_dict_robust(self.kp_extractor, checkpoint_path, device, model_name='kp_detector')
+    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'); load_state_dict_robust(self.audio_model, checkpoint_path, device, model_name='wav2vec2')
+        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'], 'auido2pose_00140-model.pth'); load_state_dict_robust(self, mapping_checkpoint, device)
+    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, jacobian, pose_coeff, expression_coeff, blink_coeff = coeff['kp_driving'], coeff['jacobian'], coeff['pose_coeff'], coeff['expression_coeff'], coeff['blink_coeff']
+        face_3d = mapping(expression_coeff, pose_coeff, blink_coeff) if blink_coeff is not None else mapping(expression_coeff, pose_coeff); 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})
+        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']
+        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.forward(np.array(pil_image)); 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'); load_state_dict_robust(self.generator, checkpoint_path, device, model_name='generator')
+    def forward(self, source_image, dense_motion, bg_param, face_3d_param=None): video_3d = self.generator(source_image, kp_driving=dense_motion, bg_param=bg_param, face_3d_param=face_3d_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'); load_state_dict_robust(self.mapping_net, checkpoint_path, device, model_name='mapping')
+        self.f_3d_mean = torch.zeros(1, 64, device=device)
+    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'); load_state_dict_robust(self.dense_motion_network, checkpoint_path, device, model_name='dense_motion')
+    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] if self.face_enhancer else x
+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 load_models():
+    checkpoint_path = './checkpoints'; config_path = './src/config'; size = 256; preprocess = 'crop'; old_version = False
+    sadtalker_instance = SadTalker(checkpoint_path, config_path, size, preprocess, old_version); print("SadTalker models loaded successfully!"); return sadtalker_instance
+if __name__ == '__main__': sadtalker_instance = load_models()

sentiment_api.py CHANGED Viewed

@@ -3,27 +3,14 @@ from main import *
 import torch
 def analyze_sentiment(text):
-    if sentiment_model is None:
-        return {"error": "Sentiment model not initialized."}
-    features = [ord(c) for c in text[:10]]
-    while len(features) < 10:
-        features.append(0)
     features_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(device)
-    with torch.no_grad():
-        output = sentiment_model(features_tensor)
-        sentiment_idx = torch.argmax(output, dim=1).item()
-        sentiment_label = "positive" if sentiment_idx == 1 else "negative"
     return {"sentiment": sentiment_label}
-def sentiment_api():
-    data = request.get_json()
-    text = data.get('text')
-    if not text:
-        return jsonify({"error": "Text is required"}), 400
     output = analyze_sentiment(text)
-    if "error" in output:
-        return jsonify({"error": output["error"]}), 500
-    return jsonify(output)

 import torch
 def analyze_sentiment(text):
+    if sentiment_model is None: return {"error": "Sentiment model not initialized."}
+    features = [ord(c) for c in text[:10]];
+    while len(features) < 10: features.append(0)
     features_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(device)
+    with torch.no_grad(): output = sentiment_model(features_tensor); sentiment_idx = torch.argmax(output, dim=1).item(); sentiment_label = "positive" if sentiment_idx == 1 else "negative"
     return {"sentiment": sentiment_label}
+def sentiment_api(text):
     output = analyze_sentiment(text)
+    if "error" in output: return {"error": output["error"]}
+    return output

stt_api.py CHANGED Viewed

@@ -1,33 +1,17 @@
-import os
-import uuid
 from flask import jsonify, send_file, request
 from main import *
-import torch
-import torchaudio
 def speech_to_text_func(audio_path):
-    if stt_model is None:
-        return {"error": "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())
-    return {"text": transcription}
-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 = speech_to_text_func(temp_audio_path)
-    os.remove(temp_audio_path)
-    if "error" in output:
-        return jsonify({"error": output["error"]}), 500
-    return jsonify(output)

+import os, uuid
 from flask import jsonify, send_file, request
 from main import *
+import torch, torchaudio
 def speech_to_text_func(audio_path):
+    if stt_model is None: return {"error": "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()); return {"text": transcription}
+def stt_api(audio_filepath):
+    output = speech_to_text_func(audio_filepath)
+    if "error" in output: return {"error": output["error"]}
+    return output

summarization_api.py CHANGED Viewed

@@ -1,27 +1,14 @@
 from flask import jsonify, send_file, request
 from main import *
-import torch
 def summarize_text(text, output_path="output_summary.txt"):
-    if summarization_model is None or summarization_tokenizer is None:
-        return "Summarization model or tokenizer not initialized."
     input_ids = summarization_tokenizer.encode(text, return_tensors="pt").to(device)
-    with torch.no_grad():
-        summary_ids = summarization_model.generate(input_ids, num_beams=4, max_length=100, early_stopping=True)
-        summary_text = summarization_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 or tokenizer not initialized.":
-        return jsonify({"error": "Summarization failed"}), 500
-    return send_file(output_file, mimetype="text/plain", as_attachment=True, download_name="output_summary.txt")

 from flask import jsonify, send_file, request
 from main import *
+import torch, io, base64
 def summarize_text(text, output_path="output_summary.txt"):
+    if summarization_model is None or summarization_tokenizer is None: return {"error": "Summarization model or tokenizer not initialized."}
     input_ids = summarization_tokenizer.encode(text, return_tensors="pt").to(device)
+    with torch.no_grad(): summary_ids = summarization_model.generate(input_ids, num_beams=4, max_length=100, early_stopping=True); summary_text = summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return {"summary_text": summary_text}
+def summarization_api(text):
+    output = summarize_text(text)
+    if "error" in output: return {"error": output["error"]}
+    return output

text_generation.py CHANGED Viewed

@@ -1,194 +1,93 @@
-import torch
-import torch.nn.functional as F
-from tqdm import trange
-import time
-from tokenxxx import *
-from main import *
-from duckduckgo_search import DDGS
-try:
-    END_OF_TEXT_TOKEN
-except NameError:
-    END_OF_TEXT_TOKEN = ""
-try:
-    SYSTEM_PROMPT
-except NameError:
-    SYSTEM_PROMPT = "Sistema: Proporcione respuestas ultra rápidas, coherentes, similares, precisas y con sentido, con razonamiento lógico y profundo."
-try:
-    MAX_XDD
-except NameError:
-    MAX_XDD = 5
-try:
-    codegen_model
-except NameError:
-    codegen_model = None
-try:
-    codegen_tokenizer
-except NameError:
-    codegen_tokenizer = None
-try:
-    summarization_model
-except NameError:
-    summarization_model = None
-try:
-    summarization_tokenizer
-except NameError:
-    summarization_tokenizer = None
-try:
-    model_gpt2
-except NameError:
-    model_gpt2 = None
-try:
-    enc
-except NameError:
-    enc = None
-try:
-    device
-except NameError:
-    device = "cpu"
-if torch.device(device).type == "cuda":
-    torch.backends.cudnn.benchmark = True
-MAX_GENERATION_LENGTH = 512
-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 _generate_sequence(model_call, context_tensor, generated, decode_fn, end_token_condition, temperature, top_k, top_p, repetition_penalty, max_length):
-    past_key_values = None
-    last_token = None
-    repetition_count = 0
-    for _ in range(max_length):
-        try:
-            outputs = model_call(context_tensor, past_key_values)
-        except Exception as e:
-            yield "<ERROR:" + str(e) + ">"
-            yield "<END_STREAM>"
-            return
-        next_token_logits = outputs[0][:, -1, :] / temperature
-        past_key_values = 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)
-        token_id = next_token.tolist()[0][0]
-        if token_id == last_token:
-            repetition_count += 1
-        else:
-            repetition_count = 0
-            last_token = token_id
-        if repetition_count >= 10:
-            yield "<END_STREAM>"
-            return
-        generated.append(token_id)
-        token_decoded = decode_fn(token_id)
-        yield token_decoded
-        if end_token_condition(token_id):
-            yield "<END_STREAM>"
-            return
-def sample_sequence(prompt, model, enc, max_length=MAX_GENERATION_LENGTH, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"):
-    context_tokens = enc.encode(prompt)
-    context_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device)
-    return _generate_sequence(
-        lambda ct, past: model(ct, past_key_values=past),
-        context_tensor,
-        list(context_tokens),
-        lambda token: enc.decode([token]),
-        lambda token: token == enc.encoder[END_OF_TEXT_TOKEN],
-        temperature, top_k, top_p, repetition_penalty, max_length
-    )
-def sample_sequence_codegen(prompt, model, tokenizer, max_length=MAX_GENERATION_LENGTH, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"):
-    context_tokens = tokenizer.encode(prompt)
-    context_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device)
-    return _generate_sequence(
-        lambda ct, past: model(input_ids=ct, past_key_values=past, labels=None),
-        context_tensor,
-        list(context_tokens),
-        lambda token: tokenizer.decode([token]),
-        lambda token: token == 50256,
-        temperature, top_k, top_p, repetition_penalty, max_length
-    )
-def summarize_text(text):
-    if summarization_model and summarization_tokenizer:
-        input_ids = summarization_tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=1024).to(device)
-        summary_ids = summarization_model.generate(
-            input_ids,
-            max_length=150,
-            min_length=40,
-            length_penalty=2.0,
-            num_beams=4,
-            early_stopping=True
-        )
-        return summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    return text[:300] + "..." if len(text) > 300 else text
-def perform_reasoning_stream(text_input, temperature, top_k, top_p, repetition_penalty, prev_context=""):
-    initial_prompt = SYSTEM_PROMPT + "\n\nUser: " + text_input + "\nAssistant:"
-    reasoning_prompt = prev_context if prev_context else initial_prompt
-    ddgs = DDGS()
-    search_results = [r for r in ddgs.text(text_input, max_results=MAX_XDD)]
-    if search_results:
-        reasoning_prompt += "\nWeb Search Results:\n"
-        for result in search_results:
-            reasoning_prompt += "- " + result['body'] + "\n"
-        reasoning_prompt += "\n"
-    if "code" in text_input.lower() or "program" in text_input.lower():
-        model_type = "code"
-    elif "summarize" in text_input.lower() or "summary" in text_input.lower():
-        model_type = "summarize"
-    elif model_gpt2 and enc:
-        model_type = "gpt2"
-    else:
-        yield "<ERROR: No se encontró un modelo adecuado>"
-        yield "<END_STREAM>"
-        return
-    if model_type == "summarize":
-        if summarization_model:
-            summary = summarize_text(text_input)
-            yield "SUMMARY_TEXT:" + summary
-            yield "<END_STREAM>"
-            return
-    accumulated_text = ""
-    current_context = reasoning_prompt
-    overlap = 256
-    while True:
-        if model_type == "code":
-            generator = sample_sequence_codegen(current_context, codegen_model, codegen_tokenizer, MAX_GENERATION_LENGTH, temperature, top_k, top_p, repetition_penalty, device)
-        elif model_type == "gpt2":
-            generator = sample_sequence(current_context, model_gpt2, enc, MAX_GENERATION_LENGTH, temperature, top_k, top_p, repetition_penalty, device)
-        chunk_text = ""
-        finished = False
-        for token in generator:
-            if token == "<END_STREAM>":
-                finished = True
-                break
-            chunk_text += token
-        if accumulated_text:
-            overlap_text = accumulated_text[-overlap:]
-            if chunk_text.startswith(overlap_text):
-                chunk_text = chunk_text[len(overlap_text):]
-        accumulated_text += chunk_text
-        yield chunk_text
-        if finished:
-            yield "<END_STREAM>"
-            break
-        current_context = accumulated_text[-overlap:] if len(accumulated_text) > overlap else accumulated_text

+import torch, torch.nn.functional as F
+from tqdm import trange
+import time
+from tokenxxx import *
+from main import *
+from duckduckgo_search import DDGS
+try: END_OF_TEXT_TOKEN
+except NameError: END_OF_TEXT_TOKEN = ""
+try: SYSTEM_PROMPT
+except NameError: SYSTEM_PROMPT = "Sistema: Proporcione respuestas ultra rápidas, coherentes, similares, precisas y con sentido, con razonamiento lógico y profundo."
+try: MAX_XDD
+except NameError: MAX_XDD = 5
+try: codegen_model
+except NameError: codegen_model = None
+try: codegen_tokenizer
+except NameError: codegen_tokenizer = None
+try: summarization_model
+except NameError: summarization_model = None
+try: summarization_tokenizer
+except NameError: summarization_tokenizer = None
+try: model_gpt2
+except NameError: model_gpt2 = None
+try: enc
+except NameError: enc = None
+try: device
+except NameError: device = "cpu"
+if torch.device(device).type == "cuda": torch.backends.cudnn.benchmark = True
+MAX_GENERATION_LENGTH = 512
+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 _generate_sequence(model_call, context_tensor, generated, decode_fn, end_token_condition, temperature, top_k, top_p, repetition_penalty, max_length):
+    past_key_values = None; last_token = None; repetition_count = 0
+    for _ in range(max_length):
+        try: outputs = model_call(context_tensor, past_key_values)
+        except Exception as e: yield "<ERROR:" + str(e) + ">"; yield "<END_STREAM>"; return
+        next_token_logits = outputs[0][:, -1, :] / temperature; past_key_values = 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)
+        token_id = next_token.tolist()[0][0]
+        if token_id == last_token: repetition_count += 1
+        else: repetition_count = 0; last_token = token_id
+        if repetition_count >= 10: yield "<END_STREAM>"; return
+        generated.append(token_id); token_decoded = decode_fn(token_id); yield token_decoded
+        if end_token_condition(token_id): yield "<END_STREAM>"; return
+def sample_sequence(prompt, model, enc, max_length=MAX_GENERATION_LENGTH, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"):
+    context_tokens = enc.encode(prompt); context_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device)
+    return _generate_sequence(lambda ct, past: model(ct, past_key_values=past), context_tensor, list(context_tokens), lambda token: enc.decode([token]), lambda token: token == enc.encoder[END_OF_TEXT_TOKEN], temperature, top_k, top_p, repetition_penalty, max_length)
+def sample_sequence_codegen(prompt, model, tokenizer, max_length=MAX_GENERATION_LENGTH, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"):
+    context_tokens = tokenizer.encode(prompt); context_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device)
+    return _generate_sequence(lambda ct, past: model(input_ids=ct, past_key_values=past, labels=None), context_tensor, list(context_tokens), lambda token: tokenizer.decode([token]), lambda token: token == 50256, temperature, top_k, top_p, repetition_penalty, max_length)
+def summarize_text(text):
+    if summarization_model and summarization_tokenizer:
+        input_ids = summarization_tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=1024).to(device); summary_ids = summarization_model.generate(input_ids, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
+        return summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return text[:300] + "..." if len(text) > 300 else text
+def perform_reasoning_stream(text_input, temperature, top_k, top_p, repetition_penalty, prev_context=""):
+    initial_prompt = SYSTEM_PROMPT + "\n\nUser: " + text_input + "\nAssistant:"; reasoning_prompt = prev_context if prev_context else initial_prompt; ddgs = DDGS()
+    search_results = [r for r in ddgs.text(text_input, max_results=MAX_XDD)];
+    if search_results: reasoning_prompt += "\nWeb Search Results:\n";
+    for result in search_results: reasoning_prompt += "- " + result['body'] + "\n"
+    reasoning_prompt += "\n"
+    if "code" in text_input.lower() or "program" in text_input.lower(): model_type = "code"
+    elif "summarize" in text_input.lower() or "summary" in text_input.lower(): model_type = "summarize"
+    elif model_gpt2 and enc: model_type = "gpt2"
+    else: yield "<ERROR: No se encontró un modelo adecuado>"; yield "<END_STREAM>"; return
+    if model_type == "summarize":
+        if summarization_model: summary = summarize_text(text_input); yield "SUMMARY_TEXT:" + summary; yield "<END_STREAM>"; return
+    accumulated_text = ""; current_context = reasoning_prompt; overlap = 256
+    while True:
+        if model_type == "code": generator = sample_sequence_codegen(current_context, codegen_model, codegen_tokenizer, MAX_GENERATION_LENGTH, temperature, top_k, top_p, repetition_penalty, device)
+        elif model_type == "gpt2": generator = sample_sequence(current_context, model_gpt2, enc, MAX_GENERATION_LENGTH, temperature, top_k, top_p, repetition_penalty, device)
+        chunk_text = ""; finished = False
+        for token in generator:
+            if token == "<END_STREAM>": finished = True; break
+            chunk_text += token
+        if accumulated_text: overlap_text = accumulated_text[-overlap:];
+        if chunk_text.startswith(overlap_text): chunk_text = chunk_text[len(overlap_text):]
+        accumulated_text += chunk_text; yield chunk_text
+        if finished: yield "<END_STREAM>"; break
+        current_context = accumulated_text[-overlap:] if len(accumulated_text) > overlap else accumulated_text

text_to_video_api.py CHANGED Viewed

@@ -1,36 +1,24 @@
-import os
-import uuid
 from flask import jsonify, send_file, request
 from main import *
-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")

+import os, uuid
 from flask import jsonify, send_file, request
 from main import *
+import torch, io
 from skimage import img_as_ubyte
+import imageio, base64
 def text_to_video_func(prompt, output_path="output_video.mp4"):
+    if text_to_video_model is None: return {"error": "Text-to-Video model not initialized."}
     video_frames_list = text_to_video_model(prompt)
+    if video_frames_list and hasattr(video_frames_list, 'frames'): export_to_video_pure(video_frames_list.frames, output_video=output_path); return output_path
+    return {"error": "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(prompt):
+    output_data = text_to_video_func(prompt)
+    if "error" in output_data: return {"error": output_data["error"]}
+    output_file = output_data;
+    with open(output_file, 'rb') as f: video_content = f.read()
+    video_base64 = base64.b64encode(video_content).decode('utf-8'); os.remove(output_file); return {"video_base64": video_base64, "mimetype": "video/mp4"}

tokenxxx.py CHANGED Viewed

@@ -1,142 +1,72 @@
-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)

+import json, re, unicodedata
 from functools import lru_cache
+import wget, 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)

translation_api.py CHANGED Viewed

@@ -2,23 +2,14 @@ from flask import jsonify, send_file, request
 from main import *
 def perform_translation(text, target_language_code='es_XX', source_language_code='en_XX'):
-    if translation_model is None:
-        return {"error": "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]
-    return {"translated_text": translation}
-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 = perform_translation(text, target_language_code=f'{target_lang}_XX', source_language_code=f'{source_lang}_XX')
-    if "error" in output:
-        return jsonify({"error": output["error"]}), 500
-    return jsonify(output)

 from main import *
 def perform_translation(text, target_language_code='es_XX', source_language_code='en_XX'):
+    if translation_model is None: return {"error": "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]; return {"translated_text": translation}
+def translation_api(text):
+    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 = perform_translation(text, target_language_code=f'{target_lang}_XX', source_language_code=f'{source_lang}_XX')
+    if "error" in output: return jsonify({"error": output["error"]}), 500
+    return output

tts_api.py CHANGED Viewed

@@ -1,26 +1,15 @@
-import os
 from flask import jsonify, send_file, request
 from main import *
-import torch
-import torchaudio
-import uuid
-def text_to_speech_func(text):
-    if tts_model is None:
-        return {"error": "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'])
-    temp_audio_path = f"temp_audio_{uuid.uuid4()}.wav"
-    torchaudio.save(temp_audio_path, audio_output.cpu(), 16000)
-    return temp_audio_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 "error" in output:
-        return jsonify({"error": output["error"]}), 500
-    return send_file(output_file, mimetype="audio/wav", as_attachment=True, download_name="output.wav")

 from flask import jsonify, send_file, request
 from main import *
+import torch, torchaudio, uuid, io, base64
+def text_to_speech_func(text, output_path="output_audio.wav"):
+    if tts_model is None: return {"error": "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(text):
     output_file = text_to_speech_func(text)
+    if isinstance(output_file, dict) and "error" in output_file: return {"error": output_file["error"]}
+    with open(output_file, 'rb') as f: audio_content = f.read()
+    audio_base64 = base64.b64encode(audio_content).decode('utf-8'); os.remove(output_file); return {"audio_base64": audio_base64, "mimetype": "audio/wav"}

xtts_api.py ADDED Viewed

	@@ -0,0 +1,21 @@

+from flask import jsonify, send_file, request
+from main import *
+import torch, torchaudio, io, base64, uuid, os
+def xtts_clone_func(text, audio_sample_path, output_path="output_xtts_audio.wav"):
+    if xtts_model is None: return {"error": "XTTS model not initialized."}
+    language = "en"; speaker_id = 0
+    try:
+        with torch.no_grad(): wav = xtts_model.inference(text=text, language_id=language, speaker_id=speaker_id, voice_sample=audio_sample_path, temperature=0.7, length_penalty=1.0)
+    except Exception as e: return {"error": f"XTTS inference failed: {e}"}
+    torchaudio.save(output_path, wav, 24000); return output_path
+def xtts_api(inputs):
+    text = inputs[0]; audio_sample_filepath = inputs[1]
+    temp_audio_path = f"temp_audio_{uuid.uuid4()}.wav"; os.rename(audio_sample_filepath, temp_audio_path)
+    output = xtts_clone_func(text, temp_audio_path); os.remove(temp_audio_path)
+    if isinstance(output, dict) and "error" in output: return {"error": output["error"]}
+    output_file = output
+    with open(output_file, 'rb') as f: audio_content = f.read()
+    audio_base64 = base64.b64encode(audio_content).decode('utf-8'); os.remove(output_file); return {"audio_base64": audio_base64, "mimetype": "audio/wav"}
+--- END OF FILE xtts_api.py ---

xxx.py CHANGED Viewed

@@ -1,142 +1,71 @@
-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)

+import json, re, unicodedata
 from functools import lru_cache
+import wget, os
 from constants import GPT2_FOLDER, ENCODER_FILE, VOCAB_FILE, END_OF_TEXT_TOKEN
 @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)