Update app_exp.py

app_exp.py

@@ -11,33 +11,21 @@ from PIL import Image
 from huggingface_hub import snapshot_download, hf_hub_download
 
 # ============================================================
-# 0️⃣ FlashAttention 3 Setup
+# 0️⃣ Install required packages
 # ============================================================
-# try:
-#     print("Attempting to download and install FlashAttention wheel...")
-#     flash_attention_wheel = hf_hub_download(
-#         repo_id="rahul7star/flash-attn-3",
-#         repo_type="model",
-#         filename="128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl",
-#     )
-#     subprocess.run(["pip", "install", flash_attention_wheel], check=True)
-#     site.addsitedir(site.getsitepackages()[0])
-#     importlib.invalidate_caches()
-#     print("✅ FlashAttention installed successfully.")
-#     enable_fa3 = True
-# except Exception as e:
-#     print(f"⚠️ Could not install FlashAttention: {e}")
-#     print("Continuing without FlashAttention...")
-#     enable_fa3 = False
+subprocess.run(["pip3", "install", "-U", "cache-dit"], check=True)
+
+
+
+import cache_dit
 
 # ============================================================
-# 1️⃣ Repository Setup
+# 1️⃣ Repository & Weights
 # ============================================================
 REPO_PATH = "LongCat-Video"
 CHECKPOINT_DIR = os.path.join(REPO_PATH, "weights", "LongCat-Video")
 
 if not os.path.exists(REPO_PATH):
-    print(f"Cloning LongCat-Video repository to '{REPO_PATH}'...")
     subprocess.run(
         ["git", "clone", "https://github.com/meituan-longcat/LongCat-Video.git", REPO_PATH],
         check=True
@@ -52,10 +40,10 @@ from longcat_video.modules.longcat_video_dit import LongCatVideoTransformer3DMod
 from longcat_video.context_parallel import context_parallel_util
 from transformers import AutoTokenizer, UMT5EncoderModel
 from diffusers.utils import export_to_video
+from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 
-# Download weights if not present
 if not os.path.exists(CHECKPOINT_DIR):
-    print(f"Downloading model weights to '{CHECKPOINT_DIR}'...")
     snapshot_download(
         repo_id="meituan-longcat/LongCat-Video",
         local_dir=CHECKPOINT_DIR,
@@ -64,33 +52,61 @@ if not os.path.exists(CHECKPOINT_DIR):
     )
 
 # ============================================================
-# 2️⃣ Device & Models
+# 2️⃣ Device & Models (with cache & quantization)
 # ============================================================
 device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.bfloat16 if device == "cuda" else torch.float32
-
-print(f"Device: {device}, dtype: {torch_dtype}")
-
+torch_dtype = torch.bfloat16 if device=="cuda" else torch.float32
 pipe = None
+
 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)
+
+    # Text encoder with 4-bit quantization
+    text_encoder = UMT5EncoderModel.from_pretrained(
+        CHECKPOINT_DIR,
+        subfolder="text_encoder",
+        torch_dtype=torch_dtype,
+        quantization_config=TransformersBitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_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)
 
+    # DiT model with FP8/4-bit quantization + cache
     dit = LongCatVideoTransformer3DModel.from_pretrained(
         CHECKPOINT_DIR,
         enable_flashattn3=enable_fa3,
-        enable_flashattn2=False,
         enable_xformers=True,
         subfolder="dit",
         cp_split_hw=cp_split_hw,
         torch_dtype=torch_dtype
     )
 
+    # Enable Cache-DiT
+    cache_dit.enable_cache(
+        cache_dit.BlockAdapter(
+            transformer=dit,
+            blocks=dit.blocks,
+            forward_pattern=cache_dit.ForwardPattern.Pattern_3,
+            check_forward_pattern=False,
+            has_separate_cfg=False
+        ),
+        cache_config=cache_dit.DBCacheConfig(
+            Fn_compute_blocks=1,
+            Bn_compute_blocks=1,
+            max_warmup_steps=5,
+            max_cached_steps=50,
+            max_continuous_cached_steps=50,
+            residual_diff_threshold=0.01,
+            num_inference_steps=50
+        )
+    )
+
     pipe = LongCatVideoPipeline(
         tokenizer=tokenizer,
         text_encoder=text_encoder,
@@ -99,55 +115,57 @@ try:
         dit=dit,
     )
     pipe.to(device)
-
-    # Load LoRA weights
-    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")
+    print("✅ Models loaded with Cache-DiT and quantization")
 
 except Exception as e:
     print(f"❌ Failed to load models: {e}")
     pipe = None
 
 # ============================================================
-# 3️⃣ Generation Helpers
+# 3️⃣ Generation Helper
 # ============================================================
 def torch_gc():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
 
-def generate_video(
+def check_duration(
     mode,
-    prompt,
-    neg_prompt,
+    prompt, 
+    neg_prompt, 
     image,
     height, width, resolution,
     seed,
     use_distill,
     use_refine,
-    duration_sec,
-    progress=gr.Progress(track_tqdm=True)
+    progress
 ):
+    if use_distill and resolution=="480p":
+        return 180
+    elif resolution=="720p":
+        return 360
+    else:
+        return 900
+
+@spaces.GPU(duration=180)
+def generate_video(mode, prompt, neg_prompt, image, height, width, resolution,
+                   seed, use_distill, use_refine, duration_sec, progress=gr.Progress(track_tqdm=True)):
+
     if pipe is None:
-        raise gr.Error("❌ Models failed to load")
+        raise gr.Error("Models not loaded")
 
-    # Adaptive FPS for faster testing
     fps = 15 if use_distill else 30
     num_frames = int(duration_sec * fps)
     generator = torch.Generator(device=device).manual_seed(int(seed))
     is_distill = use_distill or use_refine
 
-    # Stage 1
     progress(0.2, desc="Stage 1: Base Video Generation")
     pipe.dit.enable_loras(['cfg_step_lora'] if is_distill else [])
-
     num_inference_steps = 12 if is_distill else 24
     guidance_scale = 2.0 if is_distill else 4.0
     curr_neg_prompt = "" if is_distill else neg_prompt
 
-    if mode == "t2v":
+    if mode=="t2v":
         output = pipe.generate_t2v(
             prompt=prompt,
             negative_prompt=curr_neg_prompt,
@@ -176,16 +194,13 @@ def generate_video(
     pipe.dit.disable_all_loras()
     torch_gc()
 
-    # Stage 2: Optional refinement
     if use_refine:
         progress(0.5, desc="Stage 2: Refinement")
         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 = [(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
-
+        refine_image = Image.fromarray(image) if mode=='i2v' else None
         output = pipe.generate_refine(
             image=refine_image,
             prompt=prompt,
@@ -194,44 +209,43 @@ def generate_video(
             num_inference_steps=50,
             generator=generator
         )[0]
-
         pipe.dit.disable_all_loras()
         pipe.dit.disable_bsa()
         torch_gc()
 
-    # Export video
     progress(1.0, desc="Exporting video")
-    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 f:
+        export_to_video(output, f.name, fps=fps)
+        return f.name
 
 # ============================================================
 # 4️⃣ Gradio UI
 # ============================================================
-css = ".fillable{max-width: 960px !important}"
+css=".fillable{max-width:960px !important}"
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# 🎬 LongCat-Video")
+    gr.Markdown("# 🎬 LongCat-Video with Cache-DiT & Quantization")
     gr.Markdown("13.6B parameter dense video-generation model by Meituan — [[Model](https://huggingface.co/meituan-longcat/LongCat-Video)]")
 
-    with gr.Tabs() as tabs:
+    with gr.Tabs():
+        # Text-to-Video
         with gr.TabItem("Text-to-Video"):
             mode_t2v = gr.State("t2v")
             with gr.Row():
                 with gr.Column(scale=2):
                     prompt_t2v = gr.Textbox(label="Prompt", lines=4)
                     neg_prompt_t2v = gr.Textbox(label="Negative Prompt", lines=2, value="blurry, low quality")
-                    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(value=42, label="Seed")
-                    distill_t2v = gr.Checkbox(value=True, label="Use Distill Mode")
-                    refine_t2v = gr.Checkbox(value=False, label="Use Refine Mode")
-                    duration_t2v = gr.Slider(1, 20, step=1, value=2, label="Video Duration (seconds)")
-
+                    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(value=42,label="Seed")
+                    distill_t2v = gr.Checkbox(value=True,label="Use Distill Mode")
+                    refine_t2v = gr.Checkbox(value=False,label="Use Refine Mode")
+                    duration_t2v = gr.Slider(1,20,step=1,value=2,label="Video Duration (seconds)")
                     t2v_button = gr.Button("Generate Video")
                 with gr.Column(scale=3):
                     video_output_t2v = gr.Video(label="Generated Video")
 
+        # Image-to-Video
         with gr.TabItem("Image-to-Video"):
             mode_i2v = gr.State("i2v")
             with gr.Row():
@@ -240,16 +254,15 @@ with gr.Blocks(css=css) as demo:
                     prompt_i2v = gr.Textbox(label="Prompt", lines=4)
                     neg_prompt_i2v = gr.Textbox(label="Negative Prompt", lines=2, value="blurry, low quality")
                     resolution_i2v = gr.Dropdown(["480p","720p"], value="480p", label="Resolution")
-                    seed_i2v = gr.Number(value=42, label="Seed")
-                    distill_i2v = gr.Checkbox(value=True, label="Use Distill Mode")
-                    refine_i2v = gr.Checkbox(value=False, label="Use Refine Mode")
-                    duration_i2v = gr.Slider(1, 20, step=1, value=2, label="Video Duration (seconds)")
-
+                    seed_i2v = gr.Number(value=42,label="Seed")
+                    distill_i2v = gr.Checkbox(value=True,label="Use Distill Mode")
+                    refine_i2v = gr.Checkbox(value=False,label="Use Refine Mode")
+                    duration_i2v = gr.Slider(1,20,step=1,value=2,label="Video Duration (seconds)")
                     i2v_button = gr.Button("Generate Video")
                 with gr.Column(scale=3):
                     video_output_i2v = gr.Video(label="Generated Video")
 
-    # Event binding
+    # Bind events
     t2v_button.click(
         generate_video,
         inputs=[mode_t2v, prompt_t2v, neg_prompt_t2v, gr.State(None),
@@ -266,8 +279,6 @@ with gr.Blocks(css=css) as demo:
         outputs=video_output_i2v
     )
 
-# ============================================================
-# 5️⃣ Launch
-# ============================================================
-if __name__ == "__main__":
+# Launch
+if __name__=="__main__":
     demo.launch()