Update app_exp.py

app_exp.py

@@ -1,55 +1,22 @@
 import spaces
-import gradio as gr
-import torch
 import os
 import sys
-import subprocess
 import tempfile
+import datetime
 import numpy as np
-import site
-import importlib
 from PIL import Image
-from huggingface_hub import snapshot_download, hf_hub_download
-
-# ============================================================
-# 0️⃣ Install required packages
-# ============================================================
-subprocess.run(["pip3", "install", "-U", "cache-dit"], check=True)
-
-
-
-import cache_dit
-
-enable_fa3 = False  # default if FA3 cannot be loaded
-
-try:
-    print("Installing FlashAttention 3...")
-    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()
-    enable_fa3 = True
-    print("✅ FlashAttention 3 installed and enabled")
-except Exception as e:
-    print(f"⚠️ Could not install FlashAttention 3: {e}")
-    # enable_fa3 remains False
-
+import gradio as gr
+import torch
+import torch.distributed as dist
+from torchvision.io import write_video
 
 # ============================================================
-# 1️⃣ Repository & Weights
+# 1️⃣ Repo & checkpoint paths
 # ============================================================
 REPO_PATH = "LongCat-Video"
 CHECKPOINT_DIR = os.path.join(REPO_PATH, "weights", "LongCat-Video")
-
 if not os.path.exists(REPO_PATH):
-    subprocess.run(
-        ["git", "clone", "https://github.com/meituan-longcat/LongCat-Video.git", REPO_PATH],
-        check=True
-    )
+    subprocess.run(["git", "clone", "https://github.com/meituan-longcat/LongCat-Video.git", REPO_PATH], check=True)
 
 sys.path.insert(0, os.path.abspath(REPO_PATH))
 
@@ -57,248 +24,172 @@ from longcat_video.pipeline_longcat_video import LongCatVideoPipeline
 from longcat_video.modules.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 from longcat_video.modules.autoencoder_kl_wan import AutoencoderKLWan
 from longcat_video.modules.longcat_video_dit import LongCatVideoTransformer3DModel
+from longcat_video.context_parallel.context_parallel_util import init_context_parallel
 from longcat_video.context_parallel import context_parallel_util
+import cache_dit
 from transformers import AutoTokenizer, UMT5EncoderModel
-from diffusers.utils import export_to_video
-from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
-from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 
-if not os.path.exists(CHECKPOINT_DIR):
-    snapshot_download(
-        repo_id="meituan-longcat/LongCat-Video",
-        local_dir=CHECKPOINT_DIR,
-        local_dir_use_symlinks=False,
-        ignore_patterns=["*.md", "*.gitattributes", "assets/*"]
-    )
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.bfloat16 if device=="cuda" else torch.float32
+
+def torch_gc():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
 
 # ============================================================
-# 2️⃣ Device & Models (with cache & quantization)
+# 2️⃣ Model loader with cache & 4-bit/FP8 quantization
 # ============================================================
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.bfloat16 if device=="cuda" else torch.float32
-pipe = None
+def load_models(checkpoint_dir=CHECKPOINT_DIR, cp_size=1, quantize=True, cache=True):
+    cp_split_hw = context_parallel_util.get_optimal_split(cp_size)
 
-try:
-    cp_split_hw = context_parallel_util.get_optimal_split(1)
-    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR, subfolder="tokenizer", torch_dtype=torch_dtype)
+    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)
 
-    # Text encoder with 4-bit quantization
-    text_encoder = UMT5EncoderModel.from_pretrained(
-        CHECKPOINT_DIR,
-        subfolder="text_encoder",
-        torch_dtype=torch_dtype,
-        quantization_config=TransformersBitsAndBytesConfig(
+    if quantize:
+        from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
+        quant_cfg = DiffusersBitsAndBytesConfig(
             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)
+    else:
+        quant_cfg = None
 
-    # DiT model with FP8/4-bit quantization + cache
     dit = LongCatVideoTransformer3DModel.from_pretrained(
-        CHECKPOINT_DIR,
-        enable_flashattn3=enable_fa3,
-        enable_xformers=True,
+        checkpoint_dir,
         subfolder="dit",
         cp_split_hw=cp_split_hw,
-        torch_dtype=torch_dtype
+        torch_dtype=torch_dtype,
+        quantization_config=quant_cfg
     )
 
-    # 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
+    if cache:
+        from cache_dit import enable_cache, BlockAdapter, ForwardPattern, DBCacheConfig
+        enable_cache(
+            BlockAdapter(transformer=dit, blocks=dit.blocks, forward_pattern=ForwardPattern.Pattern_3),
+            cache_config=DBCacheConfig(Fn_compute_blocks=1)
         )
-    )
 
     pipe = LongCatVideoPipeline(
         tokenizer=tokenizer,
         text_encoder=text_encoder,
         vae=vae,
         scheduler=scheduler,
-        dit=dit,
+        dit=dit
     )
     pipe.to(device)
-    print("✅ Models loaded with Cache-DiT and quantization")
+    return pipe
 
-except Exception as e:
-    print(f"❌ Failed to load models: {e}")
-    pipe = None
+pipe = load_models()
 
 # ============================================================
-# 3️⃣ Generation Helper
+# 3️⃣ LoRA refinement
 # ============================================================
-def torch_gc():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
+pipe.dit.load_lora(os.path.join(CHECKPOINT_DIR, 'lora/refinement_lora.safetensors'), 'refinement_lora')
+pipe.dit.enable_loras(['refinement_lora'])
+pipe.dit.enable_bsa()
 
-def check_duration(
+# ============================================================
+# 4️⃣ Video generation function
+# ============================================================
+@spaces.GPU(duration=60)
+def generate_video(
     mode,
-    prompt, 
-    neg_prompt, 
+    prompt,
+    neg_prompt,
     image,
-    height, width, resolution,
+    height,
+    width,
+    num_frames,
     seed,
-    use_distill,
     use_refine,
-    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 not loaded")
-
-    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
-
-    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":
         output = pipe.generate_t2v(
             prompt=prompt,
-            negative_prompt=curr_neg_prompt,
+            negative_prompt=neg_prompt,
             height=height,
             width=width,
             num_frames=num_frames,
-            num_inference_steps=num_inference_steps,
-            use_distill=is_distill,
-            guidance_scale=guidance_scale,
+            num_inference_steps=50,
+            guidance_scale=4.0,
             generator=generator
         )[0]
     else:
-        pil_img = Image.fromarray(image)
+        pil_image = Image.fromarray(image)
         output = pipe.generate_i2v(
-            image=pil_img,
+            image=pil_image,
             prompt=prompt,
-            negative_prompt=curr_neg_prompt,
-            resolution=resolution,
+            negative_prompt=neg_prompt,
+            resolution=f"{height}x{width}",
             num_frames=num_frames,
-            num_inference_steps=num_inference_steps,
-            use_distill=is_distill,
-            guidance_scale=guidance_scale,
-            generator=generator
+            num_inference_steps=50,
+            guidance_scale=4.0,
+            generator=generator,
+            use_kv_cache=True,
+            offload_kv_cache=False
         )[0]
 
-    pipe.dit.disable_all_loras()
-    torch_gc()
-
     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 = [Image.fromarray(img) for img in stage1_video_pil]
-        refine_image = Image.fromarray(image) if mode=='i2v' else None
+        stage1_video_pil = [Image.fromarray(f) for f in stage1_video_pil]
+
         output = pipe.generate_refine(
-            image=refine_image,
-            prompt=prompt,
             stage1_video=stage1_video_pil,
-            num_cond_frames=1 if mode=='i2v' else 0,
+            prompt=prompt,
+            num_cond_frames=1,
             num_inference_steps=50,
             generator=generator
         )[0]
-        pipe.dit.disable_all_loras()
-        pipe.dit.disable_bsa()
-        torch_gc()
 
-    progress(1.0, desc="Exporting video")
+    output_tensor = torch.from_numpy(np.array(output))
+    output_tensor = (output_tensor*255).clamp(0,255).to(torch.uint8)
+
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
-        export_to_video(output, f.name, fps=fps)
+        write_video(f.name, output_tensor, fps=15, video_codec="libx264", options={"crf": "18"})
         return f.name
 
 # ============================================================
-# 4️⃣ Gradio UI
+# 5️⃣ Gradio interface
 # ============================================================
-css=".fillable{max-width:960px !important}"
-
-with gr.Blocks(css=css) as demo:
-    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():
-        # 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)")
-                    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():
-                with gr.Column(scale=2):
-                    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="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)")
-                    i2v_button = gr.Button("Generate Video")
-                with gr.Column(scale=3):
-                    video_output_i2v = gr.Video(label="Generated Video")
-
-    # Bind events
-    t2v_button.click(
-        generate_video,
-        inputs=[mode_t2v, prompt_t2v, neg_prompt_t2v, gr.State(None),
-                height_t2v, width_t2v, gr.State("480p"),
-                seed_t2v, distill_t2v, refine_t2v, duration_t2v],
-        outputs=video_output_t2v
-    )
-
-    i2v_button.click(
-        generate_video,
-        inputs=[mode_i2v, prompt_i2v, neg_prompt_i2v, image_i2v,
-                gr.State(None), gr.State(None), resolution_i2v,
-                seed_i2v, distill_i2v, refine_i2v, duration_i2v],
-        outputs=video_output_i2v
-    )
+with gr.Blocks() as demo:
+    gr.Markdown("# 🎬 Optimized LongCat-Video Demo (FA3 removed)")
+    with gr.Tab("Text-to-Video"):
+        prompt_t2v = gr.Textbox(label="Prompt", lines=3)
+        neg_prompt_t2v = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, low quality")
+        height_t2v = gr.Slider(256,1024,value=480,step=64,label="Height")
+        width_t2v = gr.Slider(256,1024,value=832,step=64,label="Width")
+        frames_t2v = gr.Slider(8,180,value=48,step=1,label="Frames")
+        seed_t2v = gr.Number(value=42,label="Seed",precision=0)
+        refine_t2v = gr.Checkbox(label="Use Refine",value=True)
+        out_t2v = gr.Video(label="Generated Video")
+        btn_t2v = gr.Button("Generate")
+        btn_t2v.click(
+            generate_video,
+            inputs=["t2v", prompt_t2v, neg_prompt_t2v, None, height_t2v, width_t2v, frames_t2v, seed_t2v, refine_t2v],
+            outputs=out_t2v
+        )
+    with gr.Tab("Image-to-Video"):
+        image_i2v = gr.Image(type="numpy")
+        prompt_i2v = gr.Textbox(label="Prompt", lines=3)
+        neg_prompt_i2v = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, low quality")
+        frames_i2v = gr.Slider(8,180,value=48,step=1,label="Frames")
+        seed_i2v = gr.Number(value=42,label="Seed",precision=0)
+        refine_i2v = gr.Checkbox(label="Use Refine",value=True)
+        out_i2v = gr.Video(label="Generated Video")
+        btn_i2v = gr.Button("Generate")
+        btn_i2v.click(
+            generate_video,
+            inputs=["i2v", prompt_i2v, neg_prompt_i2v, image_i2v, 480, 832, frames_i2v, seed_i2v, refine_i2v],
+            outputs=out_i2v
+        )
 
-# Launch
 if __name__=="__main__":
     demo.launch()