Update app_wip.py

app_wip.py

@@ -18,6 +18,7 @@ from utils.dataset import TextDataset
 from utils.misc import set_seed
 from demo_utils.memory import get_cuda_free_memory_gb, DynamicSwapInstaller
 
+
 # -------------------------------------------------------------------
 # Download checkpoints once when the Space starts
 # -------------------------------------------------------------------
@@ -41,6 +42,7 @@ snapshot_download(
     local_dir="./checkpoints/Reward-Forcing-T2V-1.3B",
 )
 
+
 # === Paths ===
 CONFIG_PATH = "configs/reward_forcing.yaml"
 CHECKPOINT_PATH = "checkpoints/Reward-Forcing-T2V-1.3B/rewardforcing.pt"
@@ -63,7 +65,7 @@ def reward_forcing_inference(
     Inline / simplified version of inference.py:
     - single GPU
     - text-to-video only
-    - one .txt file = N prompts, but we return only the first generated video
+    - one .txt file = N prompts, but returns only the first generated video
     """
     logs = ""
 
@@ -77,7 +79,7 @@ def reward_forcing_inference(
 
     torch.set_grad_enabled(False)
 
-    # --------------------- Stage 1: model & config init ---------------------
+    # --------------------- Phase 1: model & config init ---------------------
     progress(0.05, desc="Init: loading config")
     logs += "Loading config...\n"
     config = OmegaConf.load(CONFIG_PATH)
@@ -87,10 +89,8 @@ def reward_forcing_inference(
     progress(0.15, desc="Init: creating pipeline")
     logs += "Creating pipeline...\n"
     if hasattr(config, "denoising_step_list"):
-        # few-step sampling pipeline
         pipeline = CausalInferencePipeline(config, device=device)
     else:
-        # full diffusion pipeline
         pipeline = CausalDiffusionInferencePipeline(config, device=device)
 
     progress(0.35, desc="Init: loading checkpoint")
@@ -124,7 +124,7 @@ def reward_forcing_inference(
         dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False
     )
 
-    # --------------------- Make a clean output directory ---------------------
+    # --------------------- Clean output folder ---------------------
     progress(0.7, desc="Cleaning output folder")
     output_folder = os.path.join(
         output_root, f"rewardforcing-{num_output_frames}f", checkpoint_step
@@ -133,8 +133,7 @@ def reward_forcing_inference(
     os.makedirs(output_folder, exist_ok=True)
     logs += f"Output directory: {output_folder}\n"
 
-    # --------------------- Stage 2: inference loop ---------------------
-    # Gradio can track tqdm progress on iterable loops
+    # --------------------- Phase 2: inference loop ---------------------
     for i, batch_data in progress.tqdm(
         enumerate(dataloader),
         total=num_prompts,
@@ -153,17 +152,13 @@ def reward_forcing_inference(
 
         all_video = []
 
-        # TEXT-TO-VIDEO only (no I2V here)
+        # TEXT-TO-VIDEO only
         prompt = batch["prompts"][0]
         extended_prompt = batch.get("extended_prompts", [None])[0]
-        if extended_prompt is not None:
-            prompts = [extended_prompt]
-        else:
-            prompts = [prompt]
+        prompts = [extended_prompt] if extended_prompt else [prompt]
 
         initial_latent = None
 
-        # Noise tensor shape matches WAN2 expected latent dims
         sampled_noise = torch.randn(
             [1, num_output_frames, 16, 60, 104],
             device=device,
@@ -172,7 +167,7 @@ def reward_forcing_inference(
 
         logs += f"Generating for prompt: {prompt[:80]}...\n"
 
-        # Run WAN inference
+        # WAN2 inference
         video, latents = pipeline.inference(
             noise=sampled_noise,
             text_prompts=prompts,
@@ -183,14 +178,10 @@ def reward_forcing_inference(
 
         current_video = rearrange(video, "b t c h w -> b t h w c").cpu()
         all_video.append(current_video)
-
-        # convert to uint8 *after* concatenation
         video = 255.0 * torch.cat(all_video, dim=1)
 
-        # free VAE cache between clips
         pipeline.vae.model.clear_cache()
 
-        # Save only the first video
         if idx < num_prompts:
             model = "regular" if not use_ema else "ema"
             safe_name = prompt[:50].replace("/", "_").replace("\\", "_")
@@ -198,10 +189,10 @@ def reward_forcing_inference(
             write_video(output_path, video[0], fps=16)
             logs += f"Saved video: {output_path}\n"
 
-            progress(1.0, desc="Done ✅")
+            progress(1.0, desc="Done")
             return output_path, logs
 
-    logs += "[WARN] No video generated in loop.\n"
+    logs += "[WARN] No video generated.\n"
     return None, logs
 
 
@@ -210,18 +201,15 @@ def gradio_generate(
 ):
     """
     Triggered by Gradio:
-    - writes prompt to a temporary .txt file
-    - runs reward_forcing_inference
+    - writes prompt to a .txt file
+    - performs inference
     - returns video + logs
     """
     if not prompt or not prompt.strip():
-        raise gr.Error("Please type a text prompt 🙂")
+        raise gr.Error("Please enter a text prompt 🙂")
 
-    # Duration -> number of latent timesteps
-    if duration == "5s (21 frames)":
-        num_output_frames = 21
-    else:
-        num_output_frames = 120
+    # Duration → number of frames
+    num_output_frames = 21 if duration == "5s (21 frames)" else 120
 
     os.makedirs(PROMPT_DIR, exist_ok=True)
 
@@ -239,34 +227,37 @@ def gradio_generate(
     )
 
     if video_path is None or not os.path.exists(video_path):
-        raise gr.Error(
-            "No video generated.\n"
-            "Check the logs below for errors."
-        )
+        raise gr.Error("No video generated. Check logs for details.")
 
     return video_path, logs
 
 
 # -------------------------------------------------------------------
-# Gradio UI
+# Gradio UI — updated title + intro text
 # -------------------------------------------------------------------
 
-with gr.Blocks(title="Reward Forcing T2V Demo (inline inference)") as demo:
+with gr.Blocks(title="Reward Forcing — Text-to-Video Demo") as demo:
     gr.Markdown(
         """
-        # 🎬 Reward Forcing – Text-to-Video (inline)
+        # 🎬 Reward Forcing — Text-to-Video Demo
 
-        This version directly calls the inference logic in Python,
-        allowing Gradio to track:
-        - model initialization via `progress(...)`
-        - video generation progress via `progress.tqdm(...)`
+        Generate short videos from text prompts using a model trained with the **Reward Forcing** method.
+
+        Reward Forcing is a recent research technique that improves how well a video model follows a written description  
+        by guiding training with learned reward signals. You can learn more here:  
+        https://reward-forcing.github.io
+
+        👉 Type a prompt, click **Generate**, and the video will appear below.  
+        Longer and more detailed prompts usually produce better results.
+
+        > ⏳ The first run may take a little longer while the model loads — generation is faster afterwards.
         """
     )
 
     with gr.Row():
         prompt_in = gr.Textbox(
             label="Prompt",
-            placeholder="Ex: A cinematic shot of a spaceship flying above a neon city at night...",
+            placeholder="A cinematic shot of late-summer wheat fields moving in the wind...",
             lines=4,
         )
 
@@ -282,11 +273,7 @@ with gr.Blocks(title="Reward Forcing T2V Demo (inline inference)") as demo:
 
     with gr.Row():
         video_out = gr.Video(label="Generated Video")
-    logs_out = gr.Textbox(
-        label="Logs",
-        lines=12,
-        interactive=False,
-    )
+    logs_out = gr.Textbox(label="Logs", lines=12, interactive=False)
 
     generate_btn.click(
         fn=gradio_generate,
@@ -295,5 +282,6 @@ with gr.Blocks(title="Reward Forcing T2V Demo (inline inference)") as demo:
     )
 
 demo.queue()
+
 if __name__ == "__main__":
     demo.launch()