Upload app.py with huggingface_hub

app.py

@@ -14,7 +14,6 @@ from omegaconf import OmegaConf
 from einops import rearrange
 
 from pipeline import CausalInferencePipeline
-from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
 from wan.modules.sparse_attention import calculate_chunk_sparsities
 
 MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B"
@@ -22,7 +21,6 @@ LF_CKPT_ID = "mack-williams/Light-Forcing"
 
 # --- Model loading (module scope, eagerly on cuda) ---
 
-# Download base model if not present locally
 from huggingface_hub import snapshot_download
 
 if not os.path.exists("wan_models/Wan2.1-T2V-1.3B"):
@@ -30,10 +28,8 @@ if not os.path.exists("wan_models/Wan2.1-T2V-1.3B"):
     snapshot_download(
         repo_id=MODEL_ID,
         local_dir="wan_models/Wan2.1-T2V-1.3B",
-        local_dir_use_symlinks=False,
     )
 
-# Download Light Forcing checkpoint
 lf_ckpt_dir = "checkpoints"
 os.makedirs(lf_ckpt_dir, exist_ok=True)
 lf_ckpt_path = os.path.join(lf_ckpt_dir, "short_video_gen.pt")
@@ -51,17 +47,9 @@ config = OmegaConf.load("configs/light_forcing_short.yaml")
 default_config = OmegaConf.load("configs/default_config.yaml")
 config = OmegaConf.merge(default_config, config)
 
-# Disable efficient_deployment kernels that need sgl_kernel / TRT
-# (we use pure PyTorch fallbacks for ZeroGPU compatibility)
-model_kwargs = dict(getattr(config, "model_kwargs", {}) or {})
-efficient_deployment = model_kwargs.pop("efficient_deployment", None) or {}
-# Disable FP8 quantization, lightvae, and custom kernels
-efficient_deployment = {}
-# Remove efficient_deployment from model_kwargs to avoid duplicate kwarg
-model_kwargs.pop("efficient_deployment", None)
-
 # Calculate sparse attention sparsity schedule
 num_frame_per_block = getattr(config, "num_frame_per_block", 1)
+model_kwargs = dict(getattr(config, "model_kwargs", {}) or {})
 local_attn_size = model_kwargs.get("local_attn_size", 21)
 sparse_config = model_kwargs.get("sparse_config", {}) or {}
 NUM_OUTPUT_FRAMES = 21
@@ -70,61 +58,40 @@ sparsity_list = calculate_chunk_sparsities(
 )
 if sparsity_list:
     sparse_config["sparsity_list"] = sparsity_list
-model_kwargs["sparse_config"] = sparse_config
 
-print(f"Model kwargs: {model_kwargs}")
 print(f"Sparsity list: {sparsity_list}")
 
-# Initialize components
-text_encoder = WanTextEncoder()
-vae = WanVAEWrapper()
-transformer = WanDiffusionWrapper(
-    **model_kwargs, is_causal=True, efficient_deployment=efficient_deployment
-)
+# Initialize pipeline (CausalInferencePipeline handles all model init internally)
+pipeline = CausalInferencePipeline(config, device="cuda")
 
 # Load Light Forcing checkpoint
 state_dict = torch.load(lf_ckpt_path, map_location="cpu", weights_only=False)
-transformer.load_state_dict(state_dict["generator_ema"])
-
-text_encoder.eval()
-transformer.eval()
-vae.eval()
-
-text_encoder.requires_grad_(False)
-transformer.requires_grad_(False)
-vae.requires_grad_(False)
-
-# Move to cuda
-pipeline = CausalInferencePipeline(
-    config,
-    device="cuda",
-    generator=transformer,
-    text_encoder=text_encoder,
-    vae=vae,
-)
+pipeline.generator.load_state_dict(state_dict["generator_ema"])
 
 pipeline = pipeline.to(dtype=torch.bfloat16)
-text_encoder.to("cuda")
-transformer.to("cuda")
-vae.to("cuda")
+pipeline.text_encoder.to("cuda")
+pipeline.generator.to("cuda")
+pipeline.vae.to("cuda")
+
+pipeline.text_encoder.eval()
+pipeline.generator.eval()
+pipeline.vae.eval()
+
+pipeline.text_encoder.requires_grad_(False)
+pipeline.generator.requires_grad_(False)
+pipeline.vae.requires_grad_(False)
 
 print("Model loaded successfully!")
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=180)
 def generate(
     prompt: str,
     seed: int = 42,
     num_output_frames: int = 21,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Generate a short video from a text prompt using Light Forcing sparse attention.
-
-    Args:
-        prompt: Text description of the video to generate.
-        seed: Random seed for reproducibility.
-        num_output_frames: Number of latent frames to generate (21 ≈ 5s video at 16fps).
-    """
+    """Generate a short video from a text prompt using Light Forcing sparse attention."""
     if not prompt.strip():
         return None, "Please enter a prompt."
 
@@ -133,91 +100,23 @@ def generate(
 
     start_time = time.time()
 
-    # Initialize KV cache
-    pipeline._initialize_kv_cache(batch_size=1, dtype=torch.float16, device="cuda")
-    pipeline._initialize_crossattn_cache(batch_size=1, dtype=torch.float16, device="cuda")
-
-    # Generate noise
+    # Generate noise (bfloat16 to match model)
     noise = torch.randn(
         [1, num_output_frames, 16, 64, 96],
         device="cuda",
-        dtype=torch.float16,
+        dtype=torch.bfloat16,
     )
 
-    # Text encoding
-    conditional_dict = text_encoder(text_prompts=[prompt])
-    for key, value in conditional_dict.items():
-        conditional_dict[key] = value.to(dtype=torch.float16)
-
-    # Autoregressive block-by-block generation
-    num_blocks = num_output_frames // num_frame_per_block
-    all_num_frames = [pipeline.num_frame_per_block] * num_blocks
-    current_start_frame = 0
-    all_latents = []
-
-    for idx, current_num_frames in enumerate(all_num_frames):
-        progress((idx + 1) / len(all_num_frames), desc=f"Generating block {idx+1}/{len(all_num_frames)}")
-
-        noisy_input = noise[
-            :, current_start_frame:current_start_frame + current_num_frames
-        ]
-
-        # Denoising loop (few-step: 4 steps from denoising_step_list)
-        for index, current_timestep in enumerate(pipeline.denoising_step_list):
-            timestep = torch.ones(
-                [1, current_num_frames], device="cuda", dtype=torch.int64
-            ) * current_timestep
-
-            if index < len(pipeline.denoising_step_list) - 1:
-                _, denoised_pred = transformer(
-                    noisy_image_or_video=noisy_input,
-                    conditional_dict=conditional_dict,
-                    timestep=timestep,
-                    kv_cache=pipeline.kv_cache1,
-                    crossattn_cache=pipeline.crossattn_cache,
-                    current_start=current_start_frame * pipeline.frame_seq_length,
-                )
-                next_timestep = pipeline.denoising_step_list[index + 1]
-                noisy_input = pipeline.scheduler.add_noise(
-                    denoised_pred.flatten(0, 1),
-                    torch.randn_like(denoised_pred.flatten(0, 1)),
-                    next_timestep * torch.ones(
-                        [current_num_frames], device="cuda", dtype=torch.long
-                    ),
-                ).unflatten(0, denoised_pred.shape[:2])
-            else:
-                _, denoised_pred = transformer(
-                    noisy_image_or_video=noisy_input,
-                    conditional_dict=conditional_dict,
-                    timestep=timestep,
-                    kv_cache=pipeline.kv_cache1,
-                    crossattn_cache=pipeline.crossattn_cache,
-                    current_start=current_start_frame * pipeline.frame_seq_length,
-                )
-
-        all_latents.append(denoised_pred)
-
-        # Update KV cache with clean context
-        if idx != len(all_num_frames) - 1:
-            transformer(
-                noisy_image_or_video=denoised_pred,
-                conditional_dict=conditional_dict,
-                timestep=torch.zeros_like(timestep),
-                kv_cache=pipeline.kv_cache1,
-                crossattn_cache=pipeline.crossattn_cache,
-                current_start=current_start_frame * pipeline.frame_seq_length,
-            )
-
-        current_start_frame += current_num_frames
-
-    # Stack all latents
-    output = torch.cat(all_latents, dim=1)
-
-    # Decode to video
-    video = vae.decode_to_pixel(output, use_cache=False)
-    video = (video * 0.5 + 0.5).clamp(0, 1)
+    # Run inference using the pipeline's built-in method
+    video = pipeline.inference(
+        noise=noise,
+        text_prompts=[prompt],
+        return_latents=False,
+        profile=False,
+        low_memory=False,
+    )
 
-    # Convert to video format
+    # video: [b, t, c, h, w] in [0, 1]
     video = rearrange(video, 'b t c h w -> b t h w c').cpu()
 
     # Save as MP4
@@ -229,12 +128,11 @@ def generate(
     elapsed = time.time() - start_time
     print(f"Generation completed in {elapsed:.2f}s")
 
-    return output_path, f"Generated in {elapsed:.1f}s"
+    return output_path, f"Generated {num_output_frames} frames in {elapsed:.1f}s"
 
 
 CSS = """
 #col-container { max-width: 900px; margin: 0 auto; }
-.dark .gradio-container { color: var(--body-text-color); }
 """
 
 with gr.Blocks() as demo: