Update app_exp.py

app_exp.py

@@ -13,7 +13,7 @@ from PIL import Image
 REPO_PATH = "LongCat-Video"
 CHECKPOINT_DIR = os.path.join(REPO_PATH, "weights", "LongCat-Video")
 
-# Clone the repository if it doesn't exist
+# Clone repo if missing
 if not os.path.exists(REPO_PATH):
     print(f"Cloning LongCat-Video repository to '{REPO_PATH}'...")
     subprocess.run(
@@ -23,6 +23,7 @@ if not os.path.exists(REPO_PATH):
 
 sys.path.insert(0, os.path.abspath(REPO_PATH))
 
+# Imports from LongCat repo
 from huggingface_hub import snapshot_download
 from longcat_video.pipeline_longcat_video import LongCatVideoPipeline
 from longcat_video.modules.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
@@ -33,7 +34,7 @@ from transformers import AutoTokenizer, UMT5EncoderModel
 from diffusers.utils import export_to_video
 from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 
-# Download weights if needed
+# Download model weights if missing
 if not os.path.exists(CHECKPOINT_DIR):
     snapshot_download(
         repo_id="meituan-longcat/LongCat-Video",
@@ -49,11 +50,13 @@ torch_dtype = torch.bfloat16 if device == "cuda" else torch.float32
 print("--- Initializing Models ---")
 try:
     cp_split_hw = context_parallel_util.get_optimal_split(1)
+
     tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR, subfolder="tokenizer", torch_dtype=torch_dtype)
     text_encoder = UMT5EncoderModel.from_pretrained(CHECKPOINT_DIR, subfolder="text_encoder", torch_dtype=torch_dtype)
     vae = AutoencoderKLWan.from_pretrained(CHECKPOINT_DIR, subfolder="vae", torch_dtype=torch_dtype)
     scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(CHECKPOINT_DIR, subfolder="scheduler", torch_dtype=torch_dtype)
 
+    # ✅ 4-bit quantization enabled
     bnb_4bit_config = DiffusersBitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_quant_type="nf4",
@@ -68,7 +71,7 @@ try:
         subfolder="dit",
         cp_split_hw=cp_split_hw,
         torch_dtype=torch_dtype,
-        quantization_config=bnb_4bit_config
+        quantization_config=bnb_4bit_config  # ✅ added
     )
 
     pipe = LongCatVideoPipeline(
@@ -77,24 +80,27 @@ try:
         vae=vae,
         scheduler=scheduler,
         dit=dit,
-    )
-    pipe.to(device)
+    ).to(device)
 
     pipe.dit.load_lora(os.path.join(CHECKPOINT_DIR, 'lora/cfg_step_lora.safetensors'), 'cfg_step_lora')
     pipe.dit.load_lora(os.path.join(CHECKPOINT_DIR, 'lora/refinement_lora.safetensors'), 'refinement_lora')
 
+    print("--- Models loaded successfully ---")
 except Exception as e:
-    print("Failed to load models:", e)
+    print("❌ Model load error:", e)
     pipe = None
 
+
+# -------------------- GPU Cleanup --------------------
 def torch_gc():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
 
-def check_duration(*args, duration_t2v=2, **kwargs):
+# -------------------- Video Generation --------------------
+def check_duration(*_args, duration_t2v=2, **_kwargs):
     fps = 30
-    return duration_t2v * fps  # total frames
+    return duration_t2v * fps
 
 @spaces.GPU(duration=check_duration)
 def generate_video(
@@ -112,27 +118,25 @@ def generate_video(
     if pipe is None:
         raise gr.Error("Models failed to load.")
 
-    fps = 30
-    num_frames = duration_t2v * fps
+    generator = torch.Generator(device=device).manual_seed(int(seed))
+    num_frames = int(duration_t2v * 30)  # ✅ duration-based frame count
     print(prompt)
 
-    generator = torch.Generator(device=device).manual_seed(int(seed))
     is_distill = use_distill or use_refine
-
     if is_distill:
         pipe.dit.enable_loras(['cfg_step_lora'])
         num_inference_steps = 16
         guidance_scale = 1.0
-        current_neg_prompt = ""
+        neg = ""
     else:
         num_inference_steps = 50
         guidance_scale = 4.0
-        current_neg_prompt = neg_prompt
+        neg = neg_prompt
 
     if mode == "t2v":
         output = pipe.generate_t2v(
             prompt=prompt,
-            negative_prompt=current_neg_prompt,
+            negative_prompt=neg,
             height=height,
             width=width,
             num_frames=num_frames,
@@ -146,7 +150,7 @@ def generate_video(
         output = pipe.generate_i2v(
             image=pil_image,
             prompt=prompt,
-            negative_prompt=current_neg_prompt,
+            negative_prompt=neg,
             resolution=resolution,
             num_frames=num_frames,
             num_inference_steps=num_inference_steps,
@@ -157,80 +161,87 @@ def generate_video(
 
     if is_distill:
         pipe.dit.disable_all_loras()
+
     torch_gc()
 
     if use_refine:
-        progress(0.5, desc="Refinement")
+        progress(0.5, desc="Refining")
         pipe.dit.enable_loras(['refinement_lora'])
         pipe.dit.enable_bsa()
-        stage1_video_pil = [(frame * 255).astype(np.uint8) for frame in output]
-        stage1_video_pil = [Image.fromarray(img) for img in stage1_video_pil]
-        refine_image = Image.fromarray(image) if mode == 'i2v' else None
+
+        frames = [(frame * 255).astype(np.uint8) for frame in output]
+        frames = [Image.fromarray(f) for f in frames]
+        ref_img = Image.fromarray(image) if mode == "i2v" else None
+
         output = pipe.generate_refine(
-            image=refine_image,
+            image=ref_img,
             prompt=prompt,
-            stage1_video=stage1_video_pil,
-            num_cond_frames=1 if mode == 'i2v' else 0,
+            stage1_video=frames,
+            num_cond_frames=1 if mode == "i2v" else 0,
             num_inference_steps=50,
             generator=generator,
         )[0]
+
         pipe.dit.disable_all_loras()
         pipe.dit.disable_bsa()
         torch_gc()
 
-    progress(1.0, desc="Exporting video")
-    print("video generated")
-    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video_file:
-        export_to_video(output, temp_video_file.name, fps=fps)
-        return temp_video_file.name
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+        export_to_video(output, tmp.name, fps=30)
+        return tmp.name
 
-# --- Gradio UI ---
-css = '.fillable{max-width: 960px !important}'
+
+# -------------------- Gradio UI --------------------
+css = ".fillable{max-width:960px !important}"
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# 🎬 LongCat-Video")
     gr.Markdown("13.6B parameter dense video-generation model — [HuggingFace](https://huggingface.co/meituan-longcat/LongCat-Video)")
 
-    with gr.Tabs() as tabs:
+    with gr.Tabs():
+        # --- T2V ---
         with gr.TabItem("Text-to-Video"):
             mode_t2v = gr.State("t2v")
             prompt_t2v = gr.Textbox(label="Prompt", lines=4)
-            neg_prompt_t2v = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, low quality, static, subtitles")
-            height_t2v = gr.Slider(256, 1024, step=64, value=480, label="Height")
-            width_t2v = gr.Slider(256, 1024, step=64, value=832, label="Width")
-            seed_t2v = gr.Number(label="Seed", value=42, precision=0)
+            neg_t2v = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, low quality, static, subtitles")
+            height_t2v = gr.Slider(256, 1024, value=480, step=64, label="Height")
+            width_t2v = gr.Slider(256, 1024, value=832, step=64, label="Width")
+            seed_t2v = gr.Number(label="Seed", value=42)
             distill_t2v = gr.Checkbox(label="Use Distill Mode", value=True)
             refine_t2v = gr.Checkbox(label="Use Refine Mode", value=False)
-            duration_t2v = gr.Slider(1, 20, step=1, value=2, label="Video Duration (seconds)")
+            duration_t2v = gr.Slider(1, 20, step=1, value=2, label="Duration (seconds)")  # ✅ added
+
             t2v_button = gr.Button("Generate Video")
-            video_output_t2v = gr.Video(label="Generated Video", interactive=False)
+            video_out_t2v = gr.Video(label="Generated Video")
 
             t2v_button.click(
                 fn=generate_video,
-                inputs=[mode_t2v, prompt_t2v, neg_prompt_t2v, gr.State(None),
+                inputs=[mode_t2v, prompt_t2v, neg_t2v, gr.State(None),
                         height_t2v, width_t2v, gr.State(None),
                         seed_t2v, distill_t2v, refine_t2v, duration_t2v],
-                outputs=video_output_t2v
+                outputs=video_out_t2v
             )
 
+        # --- I2V ---
         with gr.TabItem("Image-to-Video"):
             mode_i2v = gr.State("i2v")
             image_i2v = gr.Image(type="numpy", label="Input Image")
             prompt_i2v = gr.Textbox(label="Prompt", lines=4)
-            neg_prompt_i2v = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, low quality, static, subtitles, watermark")
+            neg_i2v = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, low quality, static, subtitles, watermark")
             resolution_i2v = gr.Dropdown(["480p", "720p"], value="480p", label="Resolution")
-            seed_i2v = gr.Number(label="Seed", value=42, precision=0)
+            seed_i2v = gr.Number(label="Seed", value=42)
             distill_i2v = gr.Checkbox(label="Use Distill Mode", value=True)
             refine_i2v = gr.Checkbox(label="Use Refine Mode", value=False)
-            duration_i2v = gr.Slider(1, 20, step=1, value=2, label="Video Duration (seconds)")
+            duration_i2v = gr.Slider(1, 20, step=1, value=2, label="Duration (seconds)")  # ✅ added
+
             i2v_button = gr.Button("Generate Video")
-            video_output_i2v = gr.Video(label="Generated Video", interactive=False)
+            video_out_i2v = gr.Video(label="Generated Video")
 
             i2v_button.click(
                 fn=generate_video,
-                inputs=[mode_i2v, prompt_i2v, neg_prompt_i2v, image_i2v,
+                inputs=[mode_i2v, prompt_i2v, neg_i2v, image_i2v,
                         gr.State(None), gr.State(None), resolution_i2v,
                         seed_i2v, distill_i2v, refine_i2v, duration_i2v],
-                outputs=video_output_i2v
+                outputs=video_out_i2v
             )
 
 if __name__ == "__main__":