Update app_wip.py

app_wip.py

@@ -2,12 +2,10 @@ import os
 import sys
 import uuid
 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
@@ -79,40 +77,39 @@ def reward_forcing_inference(
 
     torch.set_grad_enabled(False)
 
-    # --------------------- BARRE 1 : init modèle / config ---------------------
-    # 4 étapes : config, pipeline, checkpoint, move to device
-    with progress.tqdm(total=4, desc="Initialisation du modèle", unit="step") as pbar:
-        logs += "Chargement de la config...\n"
-        config = OmegaConf.load(CONFIG_PATH)
-        default_config = OmegaConf.load("configs/default_config.yaml")
-        config = OmegaConf.merge(default_config, config)
-        pbar.update(1)
-
-        logs += "Initialisation de la pipeline...\n"
-        if hasattr(config, "denoising_step_list"):
-            pipeline = CausalInferencePipeline(config, device=device)
-        else:
-            pipeline = CausalDiffusionInferencePipeline(config, device=device)
-        pbar.update(1)
-
-        logs += "Chargement des poids du checkpoint...\n"
-        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]
-        pbar.update(1)
-
-        logs += "Placement du modèle sur le device...\n"
-        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)
-        pbar.update(1)
+    # --------------------- Phase 1 : init modèle / config ---------------------
+    progress(0.05, desc="Initialisation : chargement de la config")
+    logs += "Chargement de la config...\n"
+    config = OmegaConf.load(CONFIG_PATH)
+    default_config = OmegaConf.load("configs/default_config.yaml")
+    config = OmegaConf.merge(default_config, config)
+
+    progress(0.15, desc="Initialisation : création de la pipeline")
+    logs += "Initialisation de la pipeline...\n"
+    if hasattr(config, "denoising_step_list"):
+        pipeline = CausalInferencePipeline(config, device=device)
+    else:
+        pipeline = CausalDiffusionInferencePipeline(config, device=device)
+
+    progress(0.35, desc="Initialisation : chargement du checkpoint")
+    logs += "Chargement des poids du checkpoint...\n"
+    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]
+
+    progress(0.55, desc="Initialisation : placement sur le device")
+    logs += "Placement du modèle sur le device...\n"
+    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 ---------------------
+    progress(0.65, desc="Préparation du dataset")
     logs += "Préparation du dataset (TextDataset)...\n"
     dataset = TextDataset(prompt_path=prompt_txt_path, extended_prompt_path=None)
     num_prompts = len(dataset)
@@ -121,15 +118,21 @@ def reward_forcing_inference(
     from torch.utils.data import DataLoader, SequentialSampler
 
     sampler = SequentialSampler(dataset)
-    dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
+    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)
+    progress(0.7, desc="Nettoyage du dossier de sortie")
+    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"
 
-    # --------------------- BARRE 2 : boucle d'inférence ---------------------
+    # --------------------- Phase 2 : boucle d'inférence ---------------------
+    # Ici on peut utiliser progress.tqdm sur la boucle dataloader
     for i, batch_data in progress.tqdm(
         enumerate(dataloader),
         total=num_prompts,
@@ -190,13 +193,17 @@ def reward_forcing_inference(
             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"
+
+            progress(1.0, desc="Terminé ✅")
             return output_path, logs
 
     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)):
+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
@@ -219,7 +226,6 @@ def gradio_generate(prompt: str, duration: str, use_ema: bool, progress=gr.Progr
     with open(prompt_path, "w", encoding="utf-8") as f:
         f.write(prompt.strip() + "\n")
 
-    # Appel de la fonction d'inférence inline
     video_path, logs = reward_forcing_inference(
         prompt_txt_path=prompt_path,
         num_output_frames=num_output_frames,
@@ -247,9 +253,9 @@ with gr.Blocks(title="Reward Forcing T2V Demo (inline inference)") as demo:
         # 🎬 Reward Forcing – Text-to-Video (inline)
 
         Cette version appelle directement la logique d'inférence en Python,
-        ce qui permet à Gradio de suivre les `tqdm` :
-        - Initialisation du modèle
-        - Génération vidéo
+        ce qui permet à Gradio de suivre :
+        - l'initialisation du modèle (via `progress(...)`)
+        - la boucle de génération (via `progress.tqdm(...)`)
         """
     )