app.py

import torch
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.utils import export_to_video
from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_lora_to_diffusers # Keep this if it's the base
import gradio as gr
import tempfile
import os
import spaces
from huggingface_hub import hf_hub_download
import logging # For better logging

# --- Global Model Loading & LoRA Handling ---
MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
LORA_REPO_ID = "Kijai/WanVideo_comfy"
LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# This dictionary will store the manual patches extracted by the converter
MANUAL_PATCHES_STORE = {}

def _custom_convert_non_diffusers_wan_lora_to_diffusers(state_dict):
    """
    Custom converter for Wan 2.1 T2V LoRA.
    Separates LoRA A/B weights for PEFT and diff_b/diff for manual patching.
    Stores diff_b/diff in the global MANUAL_PATCHES_STORE.
    """
    global MANUAL_PATCHES_STORE
    MANUAL_PATCHES_STORE.clear() # Clear previous patches if any

    converted_state_dict_for_peft = {}
    manual_diff_patches = {}

    # Strip "diffusion_model." prefix
    original_state_dict = {
        k[len("diffusion_model.") :]: v
        for k, v in state_dict.items()
        if k.startswith("diffusion_model.")
    }

    # --- Determine number of blocks ---
    block_indices = set()
    for k_orig in original_state_dict:
        if "blocks." in k_orig:
            try:
                block_idx_str = k_orig.split("blocks.")[1].split(".")[0]
                if block_idx_str.isdigit():
                    block_indices.add(int(block_idx_str))
            except (IndexError, ValueError) as e:
                logger.warning(f"Could not parse block index from key: {k_orig} due to {e}")

    num_transformer_blocks = max(block_indices) + 1 if block_indices else 0
    if not block_indices and any("blocks." in k for k in original_state_dict):
         logger.warning("Found 'blocks.' in keys but could not determine num_transformer_blocks reliably.")


    # --- Convert Transformer Blocks (blocks.0 to blocks.N-1) ---
    for i in range(num_transformer_blocks):
        # Self-attention (attn1 in Diffusers DiT)
        for lora_key_part, diffusers_layer_name in zip(["q", "k", "v", "o"], ["to_q", "to_k", "to_v", "to_out.0"]):
            orig_lora_down_key = f"blocks.{i}.self_attn.{lora_key_part}.lora_down.weight"
            orig_lora_up_key = f"blocks.{i}.self_attn.{lora_key_part}.lora_up.weight"
            target_base_key_peft = f"blocks.{i}.attn1.{diffusers_layer_name}"
            target_base_key_manual = f"transformer.blocks.{i}.attn1.{diffusers_layer_name}"

            if orig_lora_down_key in original_state_dict and orig_lora_up_key in original_state_dict:
                converted_state_dict_for_peft[f"{target_base_key_peft}.lora_A.weight"] = original_state_dict.pop(orig_lora_down_key)
                converted_state_dict_for_peft[f"{target_base_key_peft}.lora_B.weight"] = original_state_dict.pop(orig_lora_up_key)

            orig_diff_b_key = f"blocks.{i}.self_attn.{lora_key_part}.diff_b"
            if orig_diff_b_key in original_state_dict:
                manual_diff_patches[f"{target_base_key_manual}.bias"] = original_state_dict.pop(orig_diff_b_key)

        # Cross-attention (attn2 in Diffusers DiT)
        for lora_key_part, diffusers_layer_name in zip(["q", "k", "v", "o"], ["to_q", "to_k", "to_v", "to_out.0"]):
            orig_lora_down_key = f"blocks.{i}.cross_attn.{lora_key_part}.lora_down.weight"
            orig_lora_up_key = f"blocks.{i}.cross_attn.{lora_key_part}.lora_up.weight"
            target_base_key_peft = f"blocks.{i}.attn2.{diffusers_layer_name}"
            target_base_key_manual = f"transformer.blocks.{i}.attn2.{diffusers_layer_name}"

            if orig_lora_down_key in original_state_dict and orig_lora_up_key in original_state_dict:
                converted_state_dict_for_peft[f"{target_base_key_peft}.lora_A.weight"] = original_state_dict.pop(orig_lora_down_key)
                converted_state_dict_for_peft[f"{target_base_key_peft}.lora_B.weight"] = original_state_dict.pop(orig_lora_up_key)

            orig_diff_b_key = f"blocks.{i}.cross_attn.{lora_key_part}.diff_b"
            if orig_diff_b_key in original_state_dict:
                 manual_diff_patches[f"{target_base_key_manual}.bias"] = original_state_dict.pop(orig_diff_b_key)

        # FFN
        for original_ffn_idx, diffusers_ffn_path_part in zip(["0", "2"], ["net.0.proj", "net.2"]):
            orig_lora_down_key = f"blocks.{i}.ffn.{original_ffn_idx}.lora_down.weight"
            orig_lora_up_key = f"blocks.{i}.ffn.{original_ffn_idx}.lora_up.weight"
            target_base_key_peft = f"blocks.{i}.ffn.{diffusers_ffn_path_part}"
            target_base_key_manual = f"transformer.blocks.{i}.ffn.{diffusers_ffn_path_part}"

            if orig_lora_down_key in original_state_dict and orig_lora_up_key in original_state_dict:
                converted_state_dict_for_peft[f"{target_base_key_peft}.lora_A.weight"] = original_state_dict.pop(orig_lora_down_key)
                converted_state_dict_for_peft[f"{target_base_key_peft}.lora_B.weight"] = original_state_dict.pop(orig_lora_up_key)

            orig_diff_b_key = f"blocks.{i}.ffn.{original_ffn_idx}.diff_b"
            if orig_diff_b_key in original_state_dict:
                manual_diff_patches[f"{target_base_key_manual}.bias"] = original_state_dict.pop(orig_diff_b_key)

        # Norm layers within blocks
        # LoRA has `norm3.diff` and `norm3.diff_b`. Wan2.1 base DiTBlock has `norm2`.
        norm3_diff_key = f"blocks.{i}.norm3.diff"
        norm3_diff_b_key = f"blocks.{i}.norm3.diff_b"
        target_norm_key_base_manual = f"transformer.blocks.{i}.norm2" # Diffusers DiTBlock's second norm
        if norm3_diff_key in original_state_dict:
             manual_diff_patches[f"{target_norm_key_base_manual}.weight"] = original_state_dict.pop(norm3_diff_key)
        if norm3_diff_b_key in original_state_dict:
             manual_diff_patches[f"{target_norm_key_base_manual}.bias"] = original_state_dict.pop(norm3_diff_b_key)

        # Attention QK norms
        for attn_type, diffusers_attn_block in zip(["self_attn", "cross_attn"], ["attn1", "attn2"]):
            for norm_target_suffix in ["norm_q", "norm_k"]:
                orig_norm_diff_key = f"blocks.{i}.{attn_type}.{norm_target_suffix}.diff"
                target_norm_key_manual = f"transformer.blocks.{i}.{diffusers_attn_block}.{norm_target_suffix}.weight"
                if orig_norm_diff_key in original_state_dict:
                    manual_diff_patches[target_norm_key_manual] = original_state_dict.pop(orig_norm_diff_key)

    # --- Convert Non-Block Components ---
    # Patch Embedding
    patch_emb_diff_b_key = "patch_embedding.diff_b"
    if patch_emb_diff_b_key in original_state_dict:
        manual_diff_patches["transformer.patch_embedding.bias"] = original_state_dict.pop(patch_emb_diff_b_key)

    # Text Embedding
    for orig_idx, diffusers_linear_idx in zip(["0", "2"], ["linear_1", "linear_2"]):
        orig_lora_down_key = f"text_embedding.{orig_idx}.lora_down.weight"
        orig_lora_up_key = f"text_embedding.{orig_idx}.lora_up.weight"
        target_base_key_peft = f"condition_embedder.text_embedder.{diffusers_linear_idx}"
        target_base_key_manual = f"transformer.condition_embedder.text_embedder.{diffusers_linear_idx}"
        if orig_lora_down_key in original_state_dict and orig_lora_up_key in original_state_dict:
            converted_state_dict_for_peft[f"{target_base_key_peft}.lora_A.weight"] = original_state_dict.pop(orig_lora_down_key)
            converted_state_dict_for_peft[f"{target_base_key_peft}.lora_B.weight"] = original_state_dict.pop(orig_lora_up_key)
        orig_diff_b_key = f"text_embedding.{orig_idx}.diff_b"
        if orig_diff_b_key in original_state_dict:
            manual_diff_patches[f"{target_base_key_manual}.bias"] = original_state_dict.pop(orig_diff_b_key)

    # Time Embedding
    for orig_idx, diffusers_linear_idx in zip(["0", "2"], ["linear_1", "linear_2"]):
        orig_lora_down_key = f"time_embedding.{orig_idx}.lora_down.weight"
        orig_lora_up_key = f"time_embedding.{orig_idx}.lora_up.weight"
        target_base_key_peft = f"condition_embedder.time_embedder.{diffusers_linear_idx}"
        target_base_key_manual = f"transformer.condition_embedder.time_embedder.{diffusers_linear_idx}"
        if orig_lora_down_key in original_state_dict and orig_lora_up_key in original_state_dict:
            converted_state_dict_for_peft[f"{target_base_key_peft}.lora_A.weight"] = original_state_dict.pop(orig_lora_down_key)
            converted_state_dict_for_peft[f"{target_base_key_peft}.lora_B.weight"] = original_state_dict.pop(orig_lora_up_key)
        orig_diff_b_key = f"time_embedding.{orig_idx}.diff_b"
        if orig_diff_b_key in original_state_dict:
            manual_diff_patches[f"{target_base_key_manual}.bias"] = original_state_dict.pop(orig_diff_b_key)

    # Time Projection
    orig_lora_down_key = "time_projection.1.lora_down.weight"
    orig_lora_up_key = "time_projection.1.lora_up.weight"
    target_base_key_peft = "condition_embedder.time_proj"
    target_base_key_manual = "transformer.condition_embedder.time_proj"
    if orig_lora_down_key in original_state_dict and orig_lora_up_key in original_state_dict:
        converted_state_dict_for_peft[f"{target_base_key_peft}.lora_A.weight"] = original_state_dict.pop(orig_lora_down_key)
        converted_state_dict_for_peft[f"{target_base_key_peft}.lora_B.weight"] = original_state_dict.pop(orig_lora_up_key)
    orig_diff_b_key = "time_projection.1.diff_b"
    if orig_diff_b_key in original_state_dict:
        manual_diff_patches[f"{target_base_key_manual}.bias"] = original_state_dict.pop(orig_diff_b_key)

    # Head
    orig_lora_down_key = "head.head.lora_down.weight"
    orig_lora_up_key = "head.head.lora_up.weight"
    target_base_key_peft = "proj_out" # Directly under transformer in Diffusers DiT
    target_base_key_manual = "transformer.proj_out"
    if orig_lora_down_key in original_state_dict and orig_lora_up_key in original_state_dict:
        converted_state_dict_for_peft[f"{target_base_key_peft}.lora_A.weight"] = original_state_dict.pop(orig_lora_down_key)
        converted_state_dict_for_peft[f"{target_base_key_peft}.lora_B.weight"] = original_state_dict.pop(orig_lora_up_key)
    orig_diff_b_key = "head.head.diff_b"
    if orig_diff_b_key in original_state_dict:
        manual_diff_patches[f"{target_base_key_manual}.bias"] = original_state_dict.pop(orig_diff_b_key)

    # Log any remaining keys from the original LoRA after stripping "diffusion_model."
    if len(original_state_dict) > 0:
        logger.warning(
            f"Following keys from LoRA (after stripping 'diffusion_model.') were not converted or explicitly handled for PEFT/manual patching: {original_state_dict.keys()}"
        )

    # Add "transformer." prefix for Diffusers LoraLoaderMixin to the PEFT keys
    final_peft_state_dict = {}
    for k_peft, v_peft in converted_state_dict_for_peft.items():
        final_peft_state_dict[f"transformer.{k_peft}"] = v_peft

    MANUAL_PATCHES_STORE = manual_diff_patches # Store for later use
    return final_peft_state_dict


def apply_manual_diff_patches(pipe_model, patches):
    """
    Manually applies diff_b/diff patches to the model.
    Assumes PEFT LoRA layers have already been loaded.
    """
    if not patches:
        logger.info("No manual diff patches to apply.")
        return

    logger.info(f"Applying {len(patches)} manual diff patches...")
    patched_keys_count = 0
    unpatched_keys_count = 0

    for key, diff_tensor in patches.items():
        try:
            module_to_patch = pipe_model
            attrs = key.split(".")
            
            # Navigate to the parent module
            # e.g., key = "transformer.blocks.0.attn1.to_q.bias"
            # attrs[:-1] would be ["transformer", "blocks", "0", "attn1", "to_q"]
            for attr_name in attrs[:-1]:
                if hasattr(module_to_patch, attr_name):
                    module_to_patch = getattr(module_to_patch, attr_name)
                else:
                    # If it's a PEFT wrapped layer, try to access its base_layer
                    if hasattr(module_to_patch, 'base_layer') and hasattr(module_to_patch.base_layer, attr_name):
                        module_to_patch = getattr(module_to_patch.base_layer, attr_name)
                    else:
                        raise AttributeError(f"Submodule {attr_name} not found in {module_to_patch}")
            
            param_name = attrs[-1] # "bias" or "weight"

            # Access the target layer (it might be a PEFT LoraLayer or a regular nn.Module)
            target_layer = module_to_patch

            # If PEFT wrapped it, the actual nn.Linear or nn.LayerNorm is in `base_layer`
            if hasattr(target_layer, "base_layer") and isinstance(target_layer.base_layer, (torch.nn.Linear, torch.nn.LayerNorm)):
                layer_to_modify = target_layer.base_layer
            else:
                layer_to_modify = target_layer

            if not hasattr(layer_to_modify, param_name):
                logger.error(f"Parameter '{param_name}' not found in layer '{layer_to_modify}' for key '{key}'. Skipping.")
                unpatched_keys_count +=1
                continue

            original_param = getattr(layer_to_modify, param_name)

            if original_param is None and param_name == "bias":
                # If bias is None (e.g., LayerNorm with elementwise_affine=False, or Linear(bias=False)),
                # we might need to initialize it if the diff expects to add to it.
                # For Linear layers, if bias was False, it should remain False unless LoRA intends to add one.
                # For LayerNorm, if elementwise_affine was False, adding a bias diff means it becomes affine.
                if isinstance(layer_to_modify, torch.nn.Linear):
                    if layer_to_modify.bias is None: # Check if bias was intentionally None
                        logger.warning(f"Original layer {layer_to_modify} for key '{key}' has no bias. Creating one to apply diff_b. This might be unintended if bias=False was set.")
                        layer_to_modify.bias = torch.nn.Parameter(torch.zeros_like(diff_tensor, device=diff_tensor.device, dtype=diff_tensor.dtype))
                        original_param = layer_to_modify.bias
                    else: # Should not happen if original_param was None but layer_to_modify.bias isn't
                        pass
                elif isinstance(layer_to_modify, torch.nn.LayerNorm):
                    if not layer_to_modify.elementwise_affine:
                        logger.warning(f"LayerNorm {layer_to_modify} for key '{key}' was not elementwise_affine. Applying bias diff will make it effectively affine for bias.")
                        # LayerNorm bias is initialized to zeros if elementwise_affine is True
                        layer_to_modify.bias = torch.nn.Parameter(torch.zeros_like(diff_tensor, device=diff_tensor.device, dtype=diff_tensor.dtype))
                        original_param = layer_to_modify.bias
                        # Also need to ensure weight exists if a weight diff is applied later
                        if param_name == "bias" and not hasattr(layer_to_modify, "weight"):
                             layer_to_modify.weight = torch.nn.Parameter(torch.ones_like(diff_tensor, device=diff_tensor.device, dtype=diff_tensor.dtype)) # Norm weights init to 1

            if original_param is not None:
                if original_param.shape != diff_tensor.shape:
                    logger.error(f"Shape mismatch for key '{key}': model param '{original_param.shape}', LoRA diff '{diff_tensor.shape}'. Skipping.")
                    unpatched_keys_count +=1
                    continue
                with torch.no_grad():
                    original_param.add_(diff_tensor.to(original_param.device, original_param.dtype))
                logger.info(f"Successfully applied diff to '{key}'")
                patched_keys_count +=1
            else:
                logger.warning(f"Original parameter '{param_name}' is None for key '{key}' and was not initialized. Cannot apply diff. Skipping.")
                unpatched_keys_count +=1


        except AttributeError as e:
            logger.error(f"AttributeError: Could not find module or parameter for key '{key}'. Error: {e}. Skipping.")
            unpatched_keys_count +=1
        except Exception as e:
            logger.error(f"General error applying patch for key '{key}': {e}. Skipping.")
            unpatched_keys_count +=1
    logger.info(f"Manual patching summary: {patched_keys_count} keys patched, {unpatched_keys_count} keys failed or skipped.")


# --- Model Loading ---
logger.info(f"Loading VAE for {MODEL_ID}...")
vae = AutoencoderKLWan.from_pretrained(
    MODEL_ID,
    subfolder="vae",
    torch_dtype=torch.float32 # float32 for VAE stability
)
logger.info(f"Loading Pipeline {MODEL_ID}...")
pipe = WanPipeline.from_pretrained(
    MODEL_ID,
    vae=vae,
    torch_dtype=torch.bfloat16 # bfloat16 for pipeline
)
logger.info("Moving pipeline to CUDA...")
pipe.to("cuda")

# --- LoRA Loading ---
logger.info(f"Downloading LoRA {LORA_FILENAME} from {LORA_REPO_ID}...")
causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)

logger.info("Loading LoRA weights with custom converter...")

# lora_state_dict_raw = WanPipeline.lora_state_dict(causvid_path) # This might already do some conversion

# Alternative: Load raw state_dict and then convert
from safetensors.torch import load_file as load_safetensors
raw_lora_state_dict = load_safetensors(causvid_path)

# Now call our custom converter which will populate MANUAL_PATCHES_STORE
peft_state_dict = _custom_convert_non_diffusers_wan_lora_to_diffusers(raw_lora_state_dict)

# Load the LoRA A/B matrices using PEFT
if peft_state_dict:
    pipe.load_lora_weights(
        peft_state_dict, # Pass the dictionary directly
        adapter_name="causvid_lora"
    )
    logger.info("PEFT LoRA A/B weights loaded.")
else:
    logger.warning("No PEFT-compatible LoRA weights found after conversion.")

# Apply manual diff_b and diff patches
apply_manual_diff_patches(pipe.transformer, MANUAL_PATCHES_STORE) # Apply to the transformer component
logger.info("Manual diff_b/diff patches applied.")


# --- Gradio Interface Function ---
@spaces.GPU
def generate_video(prompt, negative_prompt, height, width, num_frames, guidance_scale, steps, fps):
    logger.info("Starting video generation...")
    logger.info(f"  Prompt: {prompt}")
    logger.info(f"  Negative Prompt: {negative_prompt if negative_prompt else 'None'}")
    logger.info(f"  Height: {height}, Width: {width}")
    logger.info(f"  Num Frames: {num_frames}, FPS: {fps}")
    logger.info(f"  Guidance Scale: {guidance_scale}")

    height = (int(height) // 8) * 8
    width = (int(width) // 8) * 8
    num_frames = int(num_frames)
    fps = int(fps)

    with torch.inference_mode():
        output_frames_list = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_frames=num_frames,
            guidance_scale=float(guidance_scale),
            num_inference_steps=steps
        ).frames

    if not output_frames_list or not output_frames_list[0]:
        raise gr.Error("Model returned empty frames. Check parameters or try a different prompt.")
    output_frames = output_frames_list[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name
    export_to_video(output_frames, video_path, fps=fps)
    logger.info(f"Video successfully generated and saved to {video_path}")
    return video_path

# --- Gradio UI Definition ---
default_prompt = "A cat walks on the grass, realistic"
default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"

with gr.Blocks() as demo:
    gr.Markdown(f"""
    # Text-to-Video with Wan 2.1 (14B) + CausVid LoRA
    Powered by `diffusers` and `Wan-AI/{MODEL_ID}`.
    Model is loaded into memory when the app starts. This might take a few minutes.
    Ensure you have a GPU with sufficient VRAM (e.g., ~24GB+ for these default settings).
    """)
    # ... (rest of your Gradio UI definition remains the same) ...
    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(label="Prompt", value=default_prompt, lines=3)
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt (Optional)",
                value=default_negative_prompt,
                lines=3
            )
            with gr.Row():
                height_input = gr.Slider(minimum=256, maximum=768, step=64, value=480, label="Height (multiple of 8)")
                width_input = gr.Slider(minimum=256, maximum=1024, step=64, value=832, label="Width (multiple of 8)")
            with gr.Row():
                num_frames_input = gr.Slider(minimum=16, maximum=100, step=1, value=25, label="Number of Frames") # Reduced default for faster demo
                fps_input = gr.Slider(minimum=5, maximum=30, step=1, value=15, label="Output FPS")
            steps = gr.Slider(minimum=1.0, maximum=30.0, value=4.0, label="Steps")
            guidance_scale_input = gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=5.0, label="Guidance Scale")

            generate_button = gr.Button("Generate Video", variant="primary")

        with gr.Column(scale=3):
            video_output = gr.Video(label="Generated Video")

    generate_button.click(
        fn=generate_video,
        inputs=[
            prompt_input,
            negative_prompt_input,
            height_input,
            width_input,
            num_frames_input,
            guidance_scale_input,
            fps_input
        ],
        outputs=video_output
    )

    gr.Examples(
        examples=[
            ["A panda eating bamboo in a lush forest, cinematic lighting", default_negative_prompt, 480, 832, 25, 5.0, 4, 15],
            ["A majestic eagle soaring over snowy mountains", default_negative_prompt, 512, 768, 30, 7.0, 4, 12],
            ["Timelapse of a flower blooming, vibrant colors", "static, ugly", 384, 640, 40, 6.0, 4, 20],
            ["Astronaut walking on the moon, Earth in the background, highly detailed", default_negative_prompt, 480, 832, 20, 5.5, 4, 10],
        ],
        inputs=[prompt_input, negative_prompt_input, height_input, width_input, num_frames_input, guidance_scale_input, steps, fps_input],
        outputs=video_output,
        fn=generate_video,
        cache_examples=False
    )

if __name__ == "__main__":
    demo.queue().launch(share=True, debug=True)