Update app.py

app.py

@@ -1,4 +1,4 @@
-import os, re, math, uuid, time, shutil, logging, tempfile, threading, requests, numpy as np
+import os, re, math, uuid, time, shutil, logging, tempfile, threading, requests, asyncio, numpy as np
 from datetime import datetime, timedelta
 from collections import Counter
 
@@ -6,66 +6,80 @@ import gradio as gr
 import torch
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
 from keybert import KeyBERT
-from TTS.api import TTS
+import edge_tts
 from moviepy.editor import (
     VideoFileClip, AudioFileClip, concatenate_videoclips, concatenate_audioclips,
-    CompositeAudioClip, AudioClip, TextClip, CompositeVideoClip, VideoClip, vfx
+    CompositeAudioClip, AudioClip, TextClip, CompositeVideoClip, VideoClip
 )
 
+# ------------------- Configuración & Globals -------------------
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 logger = logging.getLogger(__name__)
 
 PEXELS_API_KEY = os.getenv("PEXELS_API_KEY")
 if not PEXELS_API_KEY:
-    raise RuntimeError("Debes definir PEXELS_API_KEY en Variables & secrets")
+    raise RuntimeError("Debes definir PEXELS_API_KEY en 'Settings' -> 'Variables & secrets'")
 
+# Carga de modelos (se hace una sola vez al iniciar el Space)
 tokenizer = GPT2Tokenizer.from_pretrained("datificate/gpt2-small-spanish")
-gpt2 = GPT2LMHeadModel.from_pretrained("datificate/gpt2-small-spanish").eval()
+gpt2_model = GPT2LMHeadModel.from_pretrained("datificate/gpt2-small-spanish").eval()
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 kw_model = KeyBERT("distilbert-base-multilingual-cased")
-tts_engine = TTS(model_name="tts_models/es/css10/vits", progress_bar=False, gpu=False)
 
 RESULTS_DIR = "video_results"
 os.makedirs(RESULTS_DIR, exist_ok=True)
-TASKS = {}
-
-# ───────── helpers ────────────────────────────────────────────────────────────────
-def gpt2_script(prompt: str, mx: int = 160) -> str:
-    ins = f"Escribe un guion corto, interesante y coherente sobre: {prompt}"
-    inp = tokenizer(ins, return_tensors="pt", truncation=True, max_length=512)
-    out = gpt2.generate(
-        **inp, max_length=mx + inp["input_ids"].shape[1], do_sample=True,
+TASKS = {} # Diccionario para almacenar el estado de las tareas
+
+# ------------------- Funciones del Pipeline de Vídeo -------------------
+def get_edge_voices_es():
+    """Obtiene y cachea la lista de voces en español de edge-tts."""
+    try:
+        voices = asyncio.run(edge_tts.list_voices())
+        es_voices = [v['ShortName'] for v in voices if v['Locale'].startswith('es-')]
+        return sorted(es_voices)
+    except Exception as e:
+        logger.error(f"No se pudieron cargar las voces de Edge TTS: {e}")
+        return ["es-ES-ElviraNeural"] # Fallback
+
+SPANISH_VOICES = get_edge_voices_es()
+
+def gpt2_script(prompt: str, max_len: int = 160) -> str:
+    instruction = f"Escribe un guion corto, interesante y coherente sobre: {prompt}"
+    inputs = tokenizer(instruction, return_tensors="pt", truncation=True, max_length=512)
+    outputs = gpt2_model.generate(
+        **inputs, max_length=max_len + inputs["input_ids"].shape[1], do_sample=True,
         top_p=0.9, top_k=40, temperature=0.7, no_repeat_ngram_size=3,
         pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
     )
-    txt = tokenizer.decode(out[0], skip_special_tokens=True)
-    return txt.split("sobre:")[-1].strip()[:mx]
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return text.split("sobre:")[-1].strip()[:max_len]
 
-def coqui_tts(text: str, path: str):
-    text = re.sub(r"[^\w\s.,!?áéíóúñÁÉÍÓÚÑ]", "", text)[:500]
-    tts_engine.tts_to_file(text=text, file_path=path)
+async def edge_tts_synth(text: str, voice: str, path: str):
+    """Sintetiza audio usando edge-tts de forma asíncrona."""
+    communicate = edge_tts.Communicate(text, voice)
+    await communicate.save(path)
 
 def keywords(text: str) -> list[str]:
-    clean = re.sub(r"[^\w\sáéíóúñÁÉÍÓÚÑ]", "", text.lower())
+    clean_text = re.sub(r"[^\w\sáéíóúñÁÉÍÓÚÑ]", "", text.lower())
     try:
-        kws = kw_model.extract_keywords(clean, stop_words="spanish", top_n=5)
+        kws = kw_model.extract_keywords(clean_text, stop_words="spanish", top_n=5)
         return [k.replace(" ", "+") for k, _ in kws if k]
     except Exception:
-        words = [w for w in clean.split() if len(w) > 4]
-        return [w for w, _ in Counter(words).most_common(5)] or ["nature"]
+        words = [w for w in clean_text.split() if len(w) > 4]
+        return [w for w, _ in Counter(words).most_common(5)] or ["naturaleza"]
 
-def pexels_search(q: str, n: int) -> list[dict]:
-    r = requests.get(
+def pexels_search(query: str, count: int) -> list[dict]:
+    res = requests.get(
         "https://api.pexels.com/videos/search",
         headers={"Authorization": PEXELS_API_KEY},
-        params={"query": q, "per_page": n, "orientation": "landscape"},
+        params={"query": query, "per_page": count, "orientation": "landscape"},
         timeout=20,
     )
-    r.raise_for_status()
-    return r.json().get("videos", [])
+    res.raise_for_status()
+    return res.json().get("videos", [])
 
-def download(url: str, folder: str) -> str | None:
+def download_file(url: str, folder: str) -> str | None:
     name = uuid.uuid4().hex + ".mp4"
     path = os.path.join(folder, name)
     with requests.get(url, stream=True, timeout=60) as r:
@@ -73,165 +87,203 @@ def download(url: str, folder: str) -> str | None:
         with open(path, "wb") as f:
             for chunk in r.iter_content(1024 * 1024):
                 f.write(chunk)
-    return path if os.path.getsize(path) > 1000 else None
+    return path if os.path.exists(path) and os.path.getsize(path) > 1000 else None
 
-def loop_audio(aclip: AudioFileClip, dur: float) -> AudioFileClip:
-    if aclip.duration >= dur:
-        return aclip.subclip(0, dur)
-    loops = math.ceil(dur / aclip.duration)
-    return concatenate_audioclips([aclip] * loops).subclip(0, dur)
+def loop_audio(audio_clip: AudioFileClip, duration: float) -> AudioFileClip:
+    if audio_clip.duration >= duration:
+        return audio_clip.subclip(0, duration)
+    loops = math.ceil(duration / audio_clip.duration)
+    return concatenate_audioclips([audio_clip] * loops).subclip(0, duration)
 
-def make_subs_clips(script: str, video_w: int, video_h: int, duration: float):
+def make_subtitle_clips(script: str, video_w: int, video_h: int, duration: float):
     sentences = [s.strip() for s in re.split(r"[.!?¿¡]", script) if s.strip()]
-    total_words = sum(len(s.split()) for s in sentences) or 1
-    word_time = duration / total_words
-    clips, cursor = [], 0.0
-    for sent in sentences:
-        n_words = len(sent.split())
-        dur = n_words * word_time
+    if not sentences: return []
+    
+    total_words = sum(len(s.split()) for s in sentences)
+    if total_words == 0: return []
+    
+    time_per_word = duration / total_words
+    clips, current_time = [], 0.0
+    
+    for sentence in sentences:
+        num_words = len(sentence.split())
+        sentence_duration = num_words * time_per_word
+        
         txt_clip = (
-            TextClip(sent, fontsize=int(video_h * 0.05), color="white",
-                     stroke_color="black", stroke_width=2, method="caption",
-                     size=(int(video_w * 0.9), None))
-            .set_start(cursor)
-            .set_duration(dur)
-            .set_position(("center", video_h * 0.85))
+            TextClip(sentence, fontsize=int(video_h * 0.05), color="white",
+                     stroke_color="black", stroke_width=1.5, method="caption",
+                     size=(int(video_w * 0.9), None), font="Arial-Bold")
+            .set_start(current_time)
+            .set_duration(sentence_duration)
+            .set_position(("center", "bottom"))
         )
         clips.append(txt_clip)
-        cursor += dur
+        current_time += sentence_duration
+        
     return clips
 
 def make_grain_clip(size: tuple[int, int], duration: float):
     w, h = size
-    def frame(_t):
-        noise = np.random.randint(0, 256, (h, w, 1), dtype=np.uint8)
+    def make_frame(t):
+        noise = np.random.randint(0, 40, (h, w, 1), dtype=np.uint8)
         return np.repeat(noise, 3, axis=2)
-    return VideoClip(frame, duration=duration).set_opacity(0.15)
-
-# ───────── video builder ──────────────────────────────────────────────────────────
-def build_video(text: str, gen_script: bool, music_fp: str | None) -> str:
-    tmp = tempfile.mkdtemp()
-    script = gpt2_script(text) if gen_script else text.strip()
-    voice_path = os.path.join(tmp, "voice.mp3")
-    coqui_tts(script, voice_path)
+    return VideoClip(make_frame, duration=duration).set_opacity(0.15)
+
+# ------------------- Función Principal de Creación de Vídeo -------------------
+def build_video(script_text: str, generate_script_flag: bool, voice: str, music_path: str | None) -> str:
+    tmp_dir = tempfile.mkdtemp()
+    
+    # 1. Guion
+    script = gpt2_script(script_text) if generate_script_flag else script_text.strip()
+    
+    # 2. Voz (TTS)
+    voice_path = os.path.join(tmp_dir, "voice.mp3")
+    asyncio.run(edge_tts_synth(script, voice, voice_path))
     voice_clip = AudioFileClip(voice_path)
-    adur = voice_clip.duration
+    video_duration = voice_clip.duration
 
-    vids = []
+    # 3. Clips de Pexels
+    video_paths = []
     for kw in keywords(script):
-        if len(vids) >= 8:
-            break
-        for v in pexels_search(kw, 2):
-            best = max(v["video_files"], key=lambda x: x["width"] * x["height"])
-            p = download(best["link"], tmp)
-            if p:
-                vids.append(p)
-            if len(vids) >= 8:
-                break
-    if not vids:
-        raise RuntimeError("Sin vídeos disponibles")
-
-    segs, acc = [], 0
-    for path in vids:
-        if acc >= adur + 2:
-            break
+        if len(video_paths) >= 8: break
+        for video_data in pexels_search(kw, 2):
+            best_file = max(video_data["video_files"], key=lambda f: f.get("width", 0) * f.get("height", 0))
+            path = download_file(best_file['link'], tmp_dir)
+            if path:
+                video_paths.append(path)
+            if len(video_paths) >= 8: break
+    if not video_paths:
+        raise RuntimeError("No se encontraron vídeos en Pexels para este guion.")
+
+    # 4. Ensamblado de vídeo base
+    segments, total_duration = [], 0
+    for path in video_paths:
+        if total_duration >= video_duration + 5: break
         clip = VideoFileClip(path)
-        seg = clip.subclip(0, min(8, clip.duration))
-        segs.append(seg)
-        acc += seg.duration
-    base = concatenate_videoclips(segs, method="chain")
-    if base.duration < adur:
-        loops = math.ceil(adur / base.duration)
-        base = concatenate_videoclips([base] * loops, method="chain")
-    base = base.subclip(0, adur)
-
-    if music_fp:
-        mclip = loop_audio(AudioFileClip(music_fp), adur).volumex(0.2)
-        audio = CompositeAudioClip([mclip, voice_clip])
-    else:
-        audio = voice_clip
+        segment = clip.subclip(0, min(8, clip.duration))
+        segments.append(segment)
+        total_duration += segment.duration
+    base_video = concatenate_videoclips(segments, method="chain")
+    if base_video.duration < video_duration:
+        base_video = loop_audio(base_video, video_duration) # Reutiliza loop_audio para vídeo si es necesario
+    base_video = base_video.subclip(0, video_duration)
 
-    subs = make_subs_clips(script, base.w, base.h, adur)
-    grain = make_grain_clip((base.w, base.h), adur)
-    final_vid = CompositeVideoClip([base, grain, *subs]).set_audio(audio)
+    # 5. Audio de fondo
+    if music_path:
+        music_clip = loop_audio(AudioFileClip(music_path), video_duration).volumex(0.20)
+        final_audio = CompositeAudioClip([music_clip, voice_clip])
+    else:
+        final_audio = voice_clip
 
-    out_path = os.path.join(tmp, "final.mp4")
-    final_vid.write_videofile(out_path, fps=24, codec="libx264", audio_codec="aac", logger=None)
-    return out_path
+    # 6. Efectos y subtítulos
+    subtitles = make_subtitle_clips(script, base_video.w, base_video.h, video_duration)
+    grain_effect = make_grain_clip(base_video.size, video_duration)
+    
+    # 7. Composición final y renderizado
+    final_video = CompositeVideoClip([base_video, grain_effect, *subtitles]).set_audio(final_audio)
+    output_path = os.path.join(tmp_dir, "final_video.mp4")
+    final_video.write_videofile(output_path, fps=24, codec="libx264", audio_codec="aac", logger=None)
+    
+    return output_path
 
-# ───────── async tasks ────────────────────────────────────────────────────────────
-def worker(tid: str, mode: str, topic: str, user_script: str, music: str | None):
+# ------------------- Sistema de Tareas Asíncronas y Limpieza -------------------
+def worker(task_id: str, mode: str, topic: str, user_script: str, voice: str, music: str | None):
     try:
-        txt = topic if mode == "Generar Guion con IA" else user_script
-        res_tmp = build_video(txt, mode == "Generar Guion con IA", music)
-        final_path = os.path.join(RESULTS_DIR, f"{tid}.mp4")
-        shutil.copy2(res_tmp, final_path)
-        TASKS[tid] = {"status": "done", "result": final_path, "ts": datetime.utcnow()}
+        text = topic if mode == "Generar Guion con IA" else user_script
+        result_tmp_path = build_video(text, mode == "Generar Guion con IA", voice, music)
+        
+        final_path = os.path.join(RESULTS_DIR, f"{task_id}.mp4")
+        shutil.copy2(result_tmp_path, final_path)
+        
+        TASKS[task_id] = {"status": "done", "result": final_path, "timestamp": datetime.utcnow()}
+        shutil.rmtree(os.path.dirname(result_tmp_path)) # Limpia el directorio temporal
     except Exception as e:
-        TASKS[tid] = {"status": "error", "error": str(e), "ts": datetime.utcnow()}
+        logger.error(f"Error en la tarea {task_id}: {e}", exc_info=True)
+        TASKS[task_id] = {"status": "error", "error": str(e), "timestamp": datetime.utcnow()}
 
-def submit(mode, topic, user_script, music):
+def submit_task(mode, topic, user_script, voice, music):
     content = topic if mode == "Generar Guion con IA" else user_script
     if not content.strip():
-        return "", "Ingresa texto"
-    tid = uuid.uuid4().hex[:8]
-    TASKS[tid] = {"status": "processing", "ts": datetime.utcnow()}
-    threading.Thread(target=worker, args=(tid, mode, topic, user_script, music), daemon=True).start()
-    return tid, f"Tarea {tid} creada"
-
-def check(tid):
-    if tid not in TASKS:
-        return None, None, "ID inválido"
-    info = TASKS[tid]
-    stat = info["status"]
-    if stat == "processing":
-        return None, None, "Procesando..."
-    if stat == "error":
-        return None, None, f"Error: {info['error']}"
-    return info["result"], info["result"], "Vídeo listo 🎉"
-
-# ───────── janitor thread ─────────────────────────────────────────────────────────
-def janitor():
+        return "", "Por favor, ingresa un tema o guion."
+    
+    task_id = uuid.uuid4().hex[:8]
+    TASKS[task_id] = {"status": "processing", "timestamp": datetime.utcnow()}
+    
+    threading.Thread(target=worker, args=(task_id, mode, topic, user_script, voice, music), daemon=True).start()
+    
+    return task_id, f"✅ Tarea creada con ID: {task_id}. Comprueba el estado en unos minutos."
+
+def check_task_status(task_id):
+    if not task_id or task_id not in TASKS:
+        return None, None, "ID de tarea no válido o no encontrado."
+    
+    task_info = TASKS[task_id]
+    status = task_info["status"]
+    
+    if status == "processing":
+        return None, None, "⏳ La tarea se está procesando..."
+    if status == "error":
+        return None, None, f"❌ Error en la tarea: {task_info['error']}"
+    if status == "done":
+        return task_info["result"], task_info["result"], "✅ ¡Vídeo listo para descargar!"
+    return None, None, "Estado desconocido."
+
+def janitor_thread():
+    """Hilo que se ejecuta periódicamente para limpiar vídeos antiguos."""
     while True:
+        time.sleep(3600)  # Cada hora
         now = datetime.utcnow()
-        for fname in os.listdir(RESULTS_DIR):
-            fpath = os.path.join(RESULTS_DIR, fname)
-            try:
-                mtime = datetime.utcfromtimestamp(os.path.getmtime(fpath))
-                if now - mtime > timedelta(hours=24):
-                    os.remove(fpath)
-                    for k, v in list(TASKS.items()):
-                        if v.get("result") == fpath:
-                            del TASKS[k]
-            except Exception:
-                pass
-        time.sleep(3600)
-
-threading.Thread(target=janitor, daemon=True).start()
-
-# ───────── gradio ui ─────────────────────────────────────────────────────────────
-with gr.Blocks(title="Generador de Vídeos IA") as demo:
+        for task_id, info in list(TASKS.items()):
+            if now - info["timestamp"] > timedelta(hours=24):
+                if info.get("result") and os.path.exists(info["result"]):
+                    try:
+                        os.remove(info["result"])
+                        logger.info(f"Limpiado vídeo antiguo: {info['result']}")
+                    except Exception as e:
+                        logger.error(f"Error al limpiar {info['result']}: {e}")
+                del TASKS[task_id]
+
+threading.Thread(target=janitor_thread, daemon=True).start()
+
+# ------------------- Interfaz de Gradio -------------------
+with gr.Blocks(title="Generador de Vídeos IA", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎬 Generador de Vídeos con IA")
+    gr.Markdown("Crea vídeos a partir de texto, con voz, música, subtítulos y efectos visuales.")
+
     with gr.Tabs():
-        with gr.TabItem("Crear Vídeo"):
-            mode = gr.Radio(["Generar Guion con IA", "Usar Mi Guion"], value="Generar Guion con IA")
-            topic = gr.Textbox(label="Tema")
-            user_script = gr.Textbox(label="Guion Completo", visible=False)
-            music = gr.Audio(type="filepath", label="Música (opcional)")
-            btn = gr.Button("Generar")
-            tid_out = gr.Textbox(label="ID de tarea")
-            msg = gr.Textbox(label="Estado")
-        with gr.TabItem("Revisar Estado"):
-            tid_in = gr.Textbox(label="ID de tarea")
-            chk = gr.Button("Verificar")
-            vid = gr.Video()
-            dlf = gr.File()
-
-    mode.change(
-        lambda m: (gr.update(visible=m == "Generar Guion con IA"), gr.update(visible=m != "Generar Guion con IA")),
-        mode, [topic, user_script]
-    )
-    btn.click(submit, [mode, topic, user_script, music], [tid_out, msg])
-    chk.click(check, tid_in, [vid, dlf, msg])
+        with gr.TabItem("1. Crear Vídeo"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    mode_radio = gr.Radio(["Generar Guion con IA", "Usar Mi Guion"], value="Generar Guion con IA", label="Elige el método")
+                    topic_textbox = gr.Textbox(label="Tema para la IA", placeholder="Ej: La historia de la Vía Láctea")
+                    script_textbox = gr.Textbox(label="Tu Guion Completo", lines=5, visible=False, placeholder="Pega aquí tu guion...")
+                    voice_dropdown = gr.Dropdown(SPANISH_VOICES, value=SPANISH_VOICES[0] if SPANISH_VOICES else None, label="Elige una voz")
+                    music_upload = gr.Audio(type="filepath", label="Música de fondo (opcional)")
+                    submit_button = gr.Button("✨ Generar Vídeo", variant="primary")
+                with gr.Column(scale=1):
+                    task_id_output = gr.Textbox(label="ID de tu Tarea (Guárdalo)", interactive=False)
+                    status_output = gr.Textbox(label="Estado", interactive=False)
+                    gr.Markdown("---")
+                    gr.Markdown("### ¿Cómo funciona?\n1. Elige un método y rellena el texto.\n2. (Opcional) Sube música de fondo.\n3. Pulsa **Generar Vídeo**.\n4. **Copia el ID** que aparecerá.\n5. Ve a la pestaña **'2. Revisar Estado'** para ver tu vídeo.")
+
+        with gr.TabItem("2. Revisar Estado"):
+            gr.Markdown("### Consulta el estado de tu vídeo")
+            with gr.Row():
+                task_id_input = gr.Textbox(label="Pega aquí el ID de tu tarea", scale=3)
+                check_button = gr.Button("🔍 Verificar", scale=1)
+            
+            status_check_output = gr.Textbox(label="Estado Actual", interactive=False)
+            video_output = gr.Video(label="Resultado del Vídeo")
+            download_file_output = gr.File(label="Descargar Fichero")
+
+    # Lógica de la interfaz
+    def toggle_textboxes(mode):
+        is_ai_mode = mode == "Generar Guion con IA"
+        return gr.update(visible=is_ai_mode), gr.update(visible=not is_ai_mode)
+
+    mode_radio.change(toggle_textboxes, inputs=mode_radio, outputs=[topic_textbox, script_textbox])
+    submit_button.click(submit_task, inputs=[mode_radio, topic_textbox, script_textbox, voice_dropdown, music_upload], outputs=[task_id_output, status_output])
+    check_button.click(check_task_status, inputs=task_id_input, outputs=[video_output, download_file_output, status_check_output])
 
 if __name__ == "__main__":
     demo.launch()