full inference integration

app_wip.py

@@ -1,40 +1,27 @@
-import sys
 import os
+import sys
 import uuid
-import subprocess
+import shutil
 from datetime import datetime
 
 import gradio as gr
+import torch
+from omegaconf import OmegaConf
+from tqdm import tqdm
+from torchvision.io import write_video
+from einops import rearrange
 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()
+from pipeline import (
+    CausalDiffusionInferencePipeline,
+    CausalInferencePipeline,
+)
+from utils.dataset import TextDataset
+from utils.misc import set_seed
+from demo_utils.memory import get_cuda_free_memory_gb, DynamicSwapInstaller
 
 # -------------------------------------------------------------------
-# Téléchargement des checkpoints (fait une fois au démarrage du Space)
+# Téléchargement des checkpoints (comme dans ton app_wip)
 # -------------------------------------------------------------------
 snapshot_download(
     repo_id="Wan-AI/Wan2.1-T2V-1.3B",
@@ -61,133 +48,215 @@ CONFIG_PATH = "configs/reward_forcing.yaml"
 CHECKPOINT_PATH = "checkpoints/Reward-Forcing-T2V-1.3B/rewardforcing.pt"
 
 PROMPT_DIR = "prompts/gradio_inputs"
-# 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)
 
 
-def run_inference(prompt: str, duration: str, use_ema: bool):
+def reward_forcing_inference(
+    prompt_txt_path: str,
+    num_output_frames: int,
+    use_ema: bool,
+    output_root: str,
+    progress: gr.Progress,
+):
     """
-    1. Écrit le prompt dans un fichier .txt
-    2. Lance inference.py avec ce fichier comme --data_path
-    3. Retourne le chemin de la vidéo .mp4 générée + les logs
+    Version simplifiée / inline d'inference.py
+    - single GPU
+    - T2V uniquement
+    - 1 prompt par fichier .txt
     """
-    import glob
+    logs = ""
+
+    # --------------------- Device & seed ---------------------
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    set_seed(0)
+
+    free_vram = get_cuda_free_memory_gb(device)
+    logs += f"Free VRAM {free_vram} GB\n"
+    low_memory = free_vram < 40
+
+    torch.set_grad_enabled(False)
+
+    # --------------------- Config & pipeline ---------------------
+    logs += "Chargement de la config...\n"
+    progress(0.05, desc="Chargement de la config")
+    config = OmegaConf.load(CONFIG_PATH)
+    default_config = OmegaConf.load("configs/default_config.yaml")
+    config = OmegaConf.merge(default_config, config)
+
+    if hasattr(config, "denoising_step_list"):
+        pipeline = CausalInferencePipeline(config, device=device)
+    else:
+        pipeline = CausalDiffusionInferencePipeline(config, device=device)
+
+    logs += "Chargement des poids du checkpoint...\n"
+    progress(0.1, desc="Chargement du checkpoint")
+    state_dict = torch.load(CHECKPOINT_PATH, map_location="cpu")
+    pipeline.generator.load_state_dict(state_dict)
+    checkpoint_step = os.path.basename(os.path.dirname(CHECKPOINT_PATH))
+    checkpoint_step = checkpoint_step.split("_")[-1]
+
+    pipeline = pipeline.to(dtype=torch.bfloat16)
+    if low_memory:
+        DynamicSwapInstaller.install_model(pipeline.text_encoder, device=device)
+    else:
+        pipeline.text_encoder.to(device=device)
+    pipeline.generator.to(device=device)
+    pipeline.vae.to(device=device)
+
+    # --------------------- Dataset / DataLoader ---------------------
+    logs += "Préparation du dataset (TextDataset)...\n"
+    progress(0.15, desc="Préparation du dataset")
+
+    dataset = TextDataset(prompt_path=prompt_txt_path, extended_prompt_path=None)
+    num_prompts = len(dataset)
+    logs += f"Number of prompts: {num_prompts}\n"
+
+    # On ne supporte que batch_size=1 ici
+    from torch.utils.data import DataLoader, SequentialSampler
+
+    sampler = SequentialSampler(dataset)
+    dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
+
+    # --------------------- Output folder (on le vide) ---------------------
+    output_folder = os.path.join(output_root, f"rewardforcing-{num_output_frames}f", checkpoint_step)
+    shutil.rmtree(output_folder, ignore_errors=True)
+    os.makedirs(output_folder, exist_ok=True)
+    logs += f"Dossier de sortie: {output_folder}\n"
+
+    progress(0.2, desc="Démarrage de l'inférence")
+
+    # --------------------- Boucle d'inférence ---------------------
+    # On tracke le tqdm de la boucle avec le Progress Gradio
+    for i, batch_data in progress.tqdm(
+        enumerate(dataloader),
+        total=num_prompts,
+        desc="Génération vidéo",
+        unit="prompt",
+    ):
+        idx = batch_data["idx"].item()
+
+        # batch_size=1 -> on simplifie
+        if isinstance(batch_data, dict):
+            batch = batch_data
+        elif isinstance(batch_data, list):
+            batch = batch_data[0]
+        else:
+            batch = batch_data
+
+        all_video = []
+
+        # TEXT-TO-VIDEO uniquement (pas d'I2V)
+        prompt = batch["prompts"][0]
+        extended_prompt = batch.get("extended_prompts", [None])[0]
+        if extended_prompt is not None:
+            prompts = [extended_prompt]
+        else:
+            prompts = [prompt]
+
+        initial_latent = None
+
+        sampled_noise = torch.randn(
+            [1, num_output_frames, 16, 60, 104],
+            device=device,
+            dtype=torch.bfloat16,
+        )
+
+        logs += f"Génération pour le prompt: {prompt[:80]}...\n"
+        progress(0.4, desc="Sampling latents")
+
+        # Appel au pipeline
+        video, latents = pipeline.inference(
+            noise=sampled_noise,
+            text_prompts=prompts,
+            return_latents=True,
+            initial_latent=initial_latent,
+            low_memory=low_memory,
+        )
+
+        progress(0.7, desc="Décodage et écriture vidéo")
+
+        current_video = rearrange(video, "b t c h w -> b t h w c").cpu()
+        all_video.append(current_video)
+
+        video = 255.0 * torch.cat(all_video, dim=1)
+
+        # Clear VAE cache
+        pipeline.vae.model.clear_cache()
+
+        # Sauvegarde vidéo
+        if idx < num_prompts:
+            model = "regular" if not use_ema else "ema"
+            # pour éviter des noms chelous, on tronque le prompt
+            safe_name = prompt[:50].replace("/", "_").replace("\\", "_")
+            output_path = os.path.join(output_folder, f"{safe_name}.mp4")
+            write_video(output_path, video[0], fps=16)
+            logs += f"Vidéo enregistrée: {output_path}\n"
+
+            # On retourne la première vidéo (une seule dans ton cas)
+            return output_path, logs
 
+    # Si on sort de la boucle sans rien (cas improbable ici)
+    logs += "[WARN] Aucune vidéo générée dans la boucle.\n"
+    return None, logs
+
+
+def gradio_generate(prompt: str, duration: str, use_ema: bool, progress=gr.Progress(track_tqdm=True)):
+    """
+    Fonction appelée par Gradio :
+    - écrit le prompt dans un .txt
+    - appelle reward_forcing_inference
+    - retourne (video_path, logs)
+    """
     if not prompt or not prompt.strip():
         raise gr.Error("Veuillez entrer un prompt texte 🙂")
 
-    # 1) Durée -> num_output_frames + dossier conforme au README
+    # Durée -> frames
     if duration == "5s (21 frames)":
         num_output_frames = 21
-        output_folder = os.path.join(OUTPUT_ROOT, "rewardforcing-5s")
-    else:  # "30s (120 frames)"
+    else:
         num_output_frames = 120
-        output_folder = os.path.join(OUTPUT_ROOT, "rewardforcing-30s")
 
-    os.makedirs(output_folder, exist_ok=True)
+    os.makedirs(PROMPT_DIR, 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 chaque ligne comme un prompt
         f.write(prompt.strip() + "\n")
 
-    # 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)
-    )
+    progress(0.01, desc="Préparation de l'inférence")
 
-    # 4) Commande inference.py
-    cmd = [
-        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",
-    ]
-    if use_ema:
-        cmd.append("--use_ema")
-
-    result = subprocess.run(
-        cmd,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        text=True,
-        cwd=cwd,  # important sur les Spaces
+    video_path, logs = reward_forcing_inference(
+        prompt_txt_path=prompt_path,
+        num_output_frames=num_output_frames,
+        use_ema=use_ema,
+        output_root=OUTPUT_ROOT,
+        progress=progress,
     )
 
-    logs = result.stdout
-    print(logs)
-
-    # 5) Si inference.py a planté, on remonte l'erreur
-    if result.returncode != 0:
+    if video_path is None or not os.path.exists(video_path):
         raise gr.Error(
-            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"
+            "Aucune vidéo trouvée après l'inférence.\n"
             "Regarde les logs ci-dessous pour voir ce qui a coincé."
         )
 
-    # 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]
-
+    progress(1.0, desc="Terminé ✅")
     return video_path, logs
 
 
-with gr.Blocks(title="Reward Forcing T2V Demo") as demo:
+# -------------------------------------------------------------------
+# UI Gradio
+# -------------------------------------------------------------------
+
+with gr.Blocks(title="Reward Forcing T2V Demo (inline inference)") as demo:
     gr.Markdown(
         """
-        # 🎬 Reward Forcing – Text-to-Video
+        # 🎬 Reward Forcing – Text-to-Video (inline)
 
-        Entrez un prompt texte, on génère un fichier `.txt` en interne
-        puis on lance `inference.py` avec ce fichier comme `--data_path`.
+        Cette version appelle directement la logique d'inférence en Python,
+        ce qui permet à Gradio de suivre le `tqdm` et d'afficher une barre de progression.
         """
     )
 
@@ -211,16 +280,17 @@ with gr.Blocks(title="Reward Forcing T2V Demo") as demo:
     with gr.Row():
         video_out = gr.Video(label="Vidéo générée")
     logs_out = gr.Textbox(
-        label="Logs de inference.py",
+        label="Logs",
         lines=12,
         interactive=False,
     )
 
     generate_btn.click(
-        fn=run_inference,
+        fn=gradio_generate,
         inputs=[prompt_in, duration, use_ema],
         outputs=[video_out, logs_out],
     )
 
+demo.queue()
 if __name__ == "__main__":
     demo.launch()