Update app.py

app.py

@@ -2,93 +2,173 @@ import gradio as gr
 import os
 from PIL import Image
 import tempfile
+from gradio_client import Client, handle_file
 import torch
 from transformers import VitsModel, AutoTokenizer
 import scipy.io.wavfile as wavfile
-from gradio_client import Client, handle_file
-import traceback
-
-# Только CPU
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-torch.set_num_threads(4)
-
-TALKING_HEAD = "Skywork/skyreels-a1-talking-head"
-model = None
-tokenizer = None
-
-def load_tts():
-    global model, tokenizer
-    if model is None:
-        print("Загружаем TTS (каз)…")
-        model = VitsModel.from_pretrained("facebook/mms-tts-kaz")
-        tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-kaz")
-        print("TTS готова")
-    return True
 
-def ru_to_kz_simple(text: str) -> str:
-    rep = {
-        "привет": "сәлем", "здравствуйте": "сәлеметсіз бе", "спасибо": "рахмет",
-        "да": "иә", "нет": "жоқ", "сегодня": "бүгін", "завтра": "ертең",
-        "урок": "сабақ", "лекция": "дәріс", "учитель": "мұғалім", "школа": "мектеп"
-    }
-    for ru, kz in rep.items():
-        text = text.replace(ru, kz).replace(ru.capitalize(), kz.capitalize())
-    return text
+# Загружаем TTS модель при старте
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
-def create_video(image: Image.Image, text: str):
-    if not image or not text.strip():
-        return None, "Загрузите фото и введите текст!"
+try:
+    tts_model = VitsModel.from_pretrained("facebook/mms-tts-rus").to(device)
+    tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
+    print("TTS модель загружена успешно!")
+except Exception as e:
+    raise RuntimeError(f"Ошибка загрузки TTS модели: {str(e)}")
 
-    load_tts()
-    text_kz = ru_to_kz_simple(text.strip())
+# Space для talking-head
+TALKING_HEAD_SPACE = "Skywork/skyreels-a1-talking-head"
 
+def inference(image: Image.Image, text: str):
+    error_msg = ""
+    video_path = None
+    audio_path = None
+    img_path = None
+    
     try:
-        # TTS
-        inputs = tokenizer(text_kz, return_tensors="pt")
+        # Валидация входных данных
+        if image is None:
+            raise ValueError("Загрузите изображение лектора!")
+        
+        if not text or not text.strip():
+            raise ValueError("Введите текст лекции!")
+        
+        if len(text) > 500:
+            raise ValueError("Текст слишком длинный! Используйте до 500 символов.")
+        
+        print(f"Генерация TTS для текста: '{text[:50]}...'")
+        
+        # Шаг 1: Генерация аудио через TTS
+        torch.manual_seed(42)
+        inputs = tts_tokenizer(text, return_tensors="pt").to(device)
+        
         with torch.no_grad():
-            waveform = model(**inputs).waveform.squeeze().cpu().numpy()
-
-        rate = model.config.sampling_rate
-        audio_path = "/tmp/audio.wav"
-        wavfile.write(audio_path, rate, (waveform * 32767).astype("int16"))
-
-        # Изображение
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-        img_path = "/tmp/img.png"
-        image.save(img_path)
-
-        # Talking head
-        client = Client(TALKING_HEAD)
+            output = tts_model(**inputs)
+            waveform = output.waveform.squeeze().cpu().numpy()
+        
+        if waveform.size == 0:
+            raise ValueError("TTS сгенерировал пустое аудио! Попробуйте другой текст.")
+        
+        # Конвертация в int16 для WAV
+        audio = (waveform * 32767).astype("int16")
+        sampling_rate = tts_model.config.sampling_rate
+        
+        # Сохранение аудио
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as audio_file:
+            wavfile.write(audio_file.name, sampling_rate, audio)
+            audio_path = audio_file.name
+        
+        print(f"TTS аудио сохранено: {audio_path} (длина: {len(waveform)/sampling_rate:.1f} сек)")
+        
+        # Шаг 2: Сохранение изображения
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as img_file:
+            # Конвертация в RGB если нужно
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            image.save(img_file.name, format='PNG')
+            img_path = img_file.name
+        
+        print(f"Изображение сохранено: {img_path}")
+        
+        # Шаг 3: Вызов talking-head API
+        print(f"Подключение к {TALKING_HEAD_SPACE}...")
+        client = Client(TALKING_HEAD_SPACE)
+        
+        # Проверяем доступные API endpoints
+        print("Доступные API методы:", client.view_api())
+        
+        # Вызов API с правильными параметрами
         result = client.predict(
-            image_path=handle_file(img_path),
-            audio_path=handle_file(audio_path),
-            guidance_scale=2.0,
-            steps=8,
+            image=handle_file(img_path),
+            audio=handle_file(audio_path),
+            guidance_scale=3.0,
+            inference_steps=10,
             api_name="/process_image_audio"
         )
-
-        video_path = result[0] if isinstance(result, (list, tuple)) else result
-        return video_path, "Бейне дайын!"
-
+        
+        print(f"Результат API: {type(result)}")
+        
+        # Обработка результата
+        if isinstance(result, tuple) and len(result) > 0:
+            video_data = result[0]
+            if isinstance(video_data, dict) and 'video' in video_data:
+                video_path = video_data['video']
+            elif isinstance(video_data, dict) and 'path' in video_data:
+                video_path = video_data['path']
+            elif isinstance(video_data, str):
+                video_path = video_data
+            else:
+                video_path = video_data
+        else:
+            video_path = result
+        
+        print(f"Видео сгенерировано: {video_path}")
+        error_msg = "✅ Видео успешно сгенерировано!"
+        
     except Exception as e:
+        error_msg = f"❌ Ошибка: {str(e)}"
+        print(f"ОШИБКА: {error_msg}")
+        import traceback
         traceback.print_exc()
-        return None, f"Қате: {e}"
-
-# === Интерфейс ===
-with gr.Blocks(title="Бейне-лектор қазақша") as app:
-    gr.Markdown("# Бейне-лектор қазақша\nФото + текст → говорящий видео-лектор")
+        
+    finally:
+        # Очистка временных файлов
+        if audio_path and os.path.exists(audio_path):
+            try:
+                os.remove(audio_path)
+                print(f"Удален временный файл: {audio_path}")
+            except:
+                pass
+        
+        if img_path and os.path.exists(img_path):
+            try:
+                os.remove(img_path)
+                print(f"Удален временный файл: {img_path}")
+            except:
+                pass
     
-    with gr.Row():
-        with gr.Column():
-            img_in = gr.Image(label="Фото лектора", type="pil")
-            txt_in = gr.Textbox(label="Текст лекции (русский)", lines=6, placeholder="Привет! Сегодня мы изучаем математику…")
-            btn = gr.Button("Сделать видео", variant="primary")
-        
-        with gr.Column():
-            video_out = gr.Video(label="Готовое видео")
-            status = gr.Textbox(label="Статус", interactive=False)
+    return video_path, error_msg
 
-    btn.click(create_video, [img_in, txt_in], [video_out, status])
+# Интерфейс Gradio
+title = "🎓 Видео-лектор с TTS (Русский)"
+description = """
+Загрузите фото лектора и введите текст лекции. 
+Система сгенерирует видео, где лектор "произносит" ваш текст!
+
+**Требования:**
+- Фото: фронтальное изображение лица
+- Текст: до 500 символов на русском языке
+"""
+
+examples = [
+    [
+        "example_image.png",
+        "Добрый день! Сегодня мы рассмотрим основы машинного обучения."
+    ]
+]
+
+iface = gr.Interface(
+    fn=inference,
+    inputs=[
+        gr.Image(type="pil", label="📸 Фото лектора"),
+        gr.Textbox(
+            lines=5, 
+            placeholder="Введите текст лекции на русском языке (до 500 символов)...", 
+            label="📝 Текст лекции"
+        )
+    ],
+    outputs=[
+        gr.Video(label="🎬 Готовое видео"),
+        gr.Textbox(label="ℹ️ Статус", interactive=False)
+    ],
+    title=title,
+    description=description,
+    flagging_mode="never",
+    examples=None,  # Добавьте примеры, если есть тестовые изображения
+    cache_examples=False
+)
 
-app.launch(server_name="0.0.0.0", server_port=7860)
+if __name__ == "__main__":
+    iface.launch()