Update app_wip.py

app_wip.py

@@ -1,63 +1,68 @@
 import sys
+import os
+import uuid
 import subprocess
+from datetime import datetime
 
-def ensure_flash_attn():
-    try:
-        import flash_attn  # noqa: F401
-        print("[init] flash-attn déjà installé")
-    except Exception as e:
-        print("[init] Installation de flash-attn (build from source)...", e, flush=True)
-        subprocess.run(
-            [
-                sys.executable,
-                "-m",
-                "pip",
-                "install",
-                "flash-attn==2.7.4.post1",
-                "--no-build-isolation",
-            ],
-            check=True,
-        )
-        import flash_attn  # noqa: F401
-        print("[init] flash-attn OK")
-
-#ensure_flash_attn()
-
+import gradio as gr
 from huggingface_hub import snapshot_download
 
+# -------------------------------------------------------------------
+# (Optionnel) flash-attn : comme tu as déjà la bonne wheel dans
+# requirements.txt, on laisse commenté pour éviter des builds lents.
+# -------------------------------------------------------------------
+# def ensure_flash_attn():
+#     try:
+#         import flash_attn  # noqa: F401
+#         print("[init] flash-attn déjà installé")
+#     except Exception as e:
+#         print("[init] Installation de flash-attn (build from source)...", e, flush=True)
+#         subprocess.run(
+#             [
+#                 sys.executable,
+#                 "-m",
+#                 "pip",
+#                 "install",
+#                 "flash-attn==2.7.4.post1",
+#                 "--no-build-isolation",
+#             ],
+#             check=True,
+#         )
+#         import flash_attn  # noqa: F401
+#         print("[init] flash-attn OK")
+
+# ensure_flash_attn()
+
+# -------------------------------------------------------------------
+# Téléchargement des checkpoints (fait une fois au démarrage du Space)
+# -------------------------------------------------------------------
 snapshot_download(
-    repo_id='Wan-AI/Wan2.1-T2V-1.3B', 
-    local_dir='./checkpoints/Wan2.1-T2V-1.3B'
+    repo_id="Wan-AI/Wan2.1-T2V-1.3B",
+    local_dir="./checkpoints/Wan2.1-T2V-1.3B",
 )
 
 snapshot_download(
-    repo_id='KlingTeam/VideoReward', 
-    local_dir='./checkpoints/Videoreward'
+    repo_id="KlingTeam/VideoReward",
+    local_dir="./checkpoints/Videoreward",
 )
 
 snapshot_download(
-    repo_id='gdhe17/Self-Forcing', 
-    local_dir='./checkpoints/ode_init.pt'
+    repo_id="gdhe17/Self-Forcing",
+    local_dir="./checkpoints/ode_init.pt",
 )
 
 snapshot_download(
-    repo_id='JaydenLu666/Reward-Forcing-T2V-1.3B', 
-    local_dir='./checkpoints/Reward-Forcing-T2V-1.3B'
+    repo_id="JaydenLu666/Reward-Forcing-T2V-1.3B",
+    local_dir="./checkpoints/Reward-Forcing-T2V-1.3B",
 )
 
-import os
-import uuid
-import subprocess
-from datetime import datetime
-
-import gradio as gr
-
-# === Chemins à adapter si besoin ===
+# === Chemins ===
 CONFIG_PATH = "configs/reward_forcing.yaml"
 CHECKPOINT_PATH = "checkpoints/Reward-Forcing-T2V-1.3B/rewardforcing.pt"
 
 PROMPT_DIR = "prompts/gradio_inputs"
-OUTPUT_ROOT = "videos/gradio_outputs"
+# on garde OUTPUT_ROOT mais on va aussi coller au README pour l'output
+OUTPUT_ROOT = "videos"
 
 os.makedirs(PROMPT_DIR, exist_ok=True)
 os.makedirs(OUTPUT_ROOT, exist_ok=True)
@@ -67,40 +72,54 @@ def run_inference(prompt: str, duration: str, use_ema: bool):
     """
     1. Écrit le prompt dans un fichier .txt
     2. Lance inference.py avec ce fichier comme --data_path
-    3. Retourne le chemin de la première vidéo .mp4 générée + les logs
+    3. Retourne le chemin de la vidéo .mp4 générée + les logs
     """
+    import glob
+
     if not prompt or not prompt.strip():
         raise gr.Error("Veuillez entrer un prompt texte 🙂")
 
-    # 1) On mappe la durée choisie → num_output_frames
+    # 1) Durée -> num_output_frames + dossier conforme au README
     if duration == "5s (21 frames)":
         num_output_frames = 21
+        output_folder = os.path.join(OUTPUT_ROOT, "rewardforcing-5s")
     else:  # "30s (120 frames)"
         num_output_frames = 120
+        output_folder = os.path.join(OUTPUT_ROOT, "rewardforcing-30s")
+
+    os.makedirs(output_folder, exist_ok=True)
 
     # 2) Fichier .txt temporaire pour le prompt
     prompt_id = uuid.uuid4().hex[:8]
     prompt_path = os.path.join(PROMPT_DIR, f"prompt_{prompt_id}.txt")
 
     with open(prompt_path, "w", encoding="utf-8") as f:
-        # TextDataset lit juste chaque ligne comme un prompt
+        # TextDataset lit chaque ligne comme un prompt
         f.write(prompt.strip() + "\n")
 
-    # 3) Dossier de sortie unique pour cette génération
-    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-    output_folder = os.path.join(OUTPUT_ROOT, f"{ts}_{prompt_id}")
-    os.makedirs(output_folder, exist_ok=True)
+    # 3) On sauve la liste des vidéos AVANT l'inférence
+    cwd = os.path.dirname(os.path.abspath(__file__))
+    before_mp4s = set(
+        os.path.relpath(p, cwd)
+        for p in glob.glob(os.path.join(cwd, "videos", "**", "*.mp4"), recursive=True)
+    )
 
     # 4) Commande inference.py
     cmd = [
-        "python",
+        sys.executable,
         "inference.py",
-        "--num_output_frames", str(num_output_frames),
-        "--config_path", CONFIG_PATH,
-        "--checkpoint_path", CHECKPOINT_PATH,
-        "--output_folder", output_folder,
-        "--data_path", prompt_path,
-        "--num_samples", "1",
+        "--num_output_frames",
+        str(num_output_frames),
+        "--config_path",
+        CONFIG_PATH,
+        "--checkpoint_path",
+        CHECKPOINT_PATH,
+        "--output_folder",
+        output_folder,
+        "--data_path",
+        prompt_path,
+        "--num_samples",
+        "1",
     ]
     if use_ema:
         cmd.append("--use_ema")
@@ -110,21 +129,55 @@ def run_inference(prompt: str, duration: str, use_ema: bool):
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
         text=True,
+        cwd=cwd,  # important sur les Spaces
     )
 
     logs = result.stdout
     print(logs)
 
-    # 5) On récupère la première vidéo produite
-    mp4s = [f for f in os.listdir(output_folder) if f.lower().endswith(".mp4")]
-    if not mp4s:
+    # 5) Si inference.py a planté, on remonte l'erreur
+    if result.returncode != 0:
         raise gr.Error(
-            "Aucune vidéo trouvée dans le dossier de sortie.\n"
+            f"inference.py a retourné un code d'erreur ({result.returncode}).\n\n"
+            "Regarde les logs ci-dessous pour les détails."
+        )
+
+    # 6) On regarde les vidéos APRÈS l'inférence
+    after_mp4s_abs = glob.glob(os.path.join(cwd, "videos", "**", "*.mp4"), recursive=True)
+    after_mp4s = set(os.path.relpath(p, cwd) for p in after_mp4s_abs)
+
+    new_mp4s = list(after_mp4s - before_mp4s)
+
+    # Debug : log de tout ce qui a été trouvé
+    logs += "\n\n[DEBUG] Fichiers .mp4 AVANT:\n"
+    logs += "\n".join(sorted(before_mp4s)) if before_mp4s else "[aucun]\n"
+    logs += "\n\n[DEBUG] Fichiers .mp4 APRÈS:\n"
+    logs += "\n".join(sorted(after_mp4s)) if after_mp4s else "[aucun]\n"
+
+    if not new_mp4s:
+        # Pas de nouvelle vidéo détectée. En dernier recours,
+        # on prend la plus récente dans tout `videos/` si elle existe.
+        if after_mp4s_abs:
+            after_mp4s_abs.sort(key=os.path.getmtime, reverse=True)
+            fallback_video = after_mp4s_abs[0]
+            logs += (
+                "\n\n[WARN] Aucune nouvelle vidéo détectée, "
+                "on utilise la plus récente trouvée: "
+                f"{os.path.relpath(fallback_video, cwd)}"
+            )
+            return fallback_video, logs
+
+        # Vraiment aucune vidéo
+        raise gr.Error(
+            "Aucune vidéo .mp4 trouvée dans le dossier de sortie.\n"
             "Regarde les logs ci-dessous pour voir ce qui a coincé."
         )
 
-    mp4s.sort()
-    video_path = os.path.join(output_folder, mp4s[0])
+    # On prend la nouvelle vidéo la plus récente
+    new_mp4s_abs = [os.path.join(cwd, p) for p in new_mp4s]
+    new_mp4s_abs.sort(key=os.path.getmtime, reverse=True)
+    video_path = new_mp4s_abs[0]
+
     return video_path, logs
 
 
@@ -159,7 +212,7 @@ with gr.Blocks(title="Reward Forcing T2V Demo") as demo:
         video_out = gr.Video(label="Vidéo générée")
     logs_out = gr.Textbox(
         label="Logs de inference.py",
-        lines=10,
+        lines=12,
         interactive=False,
     )