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
    skipped_keys_details = []

    for key, diff_tensor in patches.items():
        try:
            # key is like "transformer.blocks.0.attn1.to_q.bias"
            current_module = pipe_model # Starts from pipe.transformer
            path_parts = key.split('.')[1:] # Remove "transformer." prefix for getattr navigation
                                           # e.g., ["blocks", "0", "attn1", "to_q", "bias"]
            
            # Navigate to the parent module of the parameter
            # Example: for "blocks.0.attn1.to_q.bias", parent_module_path is "blocks.0.attn1.to_q"
            parent_module_path = path_parts[:-1]
            param_name_to_patch = path_parts[-1] # "bias" or "weight"

            for part in parent_module_path:
                if hasattr(current_module, part):
                    current_module = getattr(current_module, part)
                elif hasattr(current_module, 'base_layer') and hasattr(current_module.base_layer, part):
                    # This case is unlikely here as we are navigating *to* the layer,
                    # not trying to access a sub-component of a base_layer.
                    # PEFT wrapping affects the layer itself, not its parent structure.
                    current_module = getattr(current_module.base_layer, part)
                else:
                    raise AttributeError(f"Submodule '{part}' not found in path '{'.'.join(parent_module_path)}' within {key}")

            # Now, current_module is the layer whose parameter we want to patch
            # e.g., if key was transformer.blocks.0.attn1.to_q.bias,
            # current_module is the to_q Linear layer (or LoraLayer wrapping it)

            layer_to_modify = current_module
            # If PEFT wrapped the Linear layer (common for attention q,k,v,o and ffn projections)
            if hasattr(layer_to_modify, "base_layer") and isinstance(layer_to_modify.base_layer, (torch.nn.Linear, torch.nn.LayerNorm)):
                actual_param_owner = layer_to_modify.base_layer
            else: # For non-wrapped layers like LayerNorm, or if it's already the base_layer
                actual_param_owner = layer_to_modify
            
            if not hasattr(actual_param_owner, param_name_to_patch):
                skipped_keys_details.append(f"Key: {key}, Reason: Parameter '{param_name_to_patch}' not found in layer '{actual_param_owner}'. Layer type: {type(actual_param_owner)}")
                unpatched_keys_count += 1
                continue

            original_param = getattr(actual_param_owner, param_name_to_patch)

            if original_param is None and param_name_to_patch == "bias":
                logger.info(f"Key '{key}': Original bias is None. Attempting to initialize.")
                if isinstance(actual_param_owner, torch.nn.Linear) or isinstance(actual_param_owner, torch.nn.LayerNorm):
                    # For LayerNorm, bias exists if elementwise_affine=True (default).
                    # If it was False, we are making it affine by adding a bias.
                    # For Linear, if bias was False, we are adding one.
                    actual_param_owner.bias = torch.nn.Parameter(torch.zeros_like(diff_tensor, device=diff_tensor.device, dtype=diff_tensor.dtype))
                    original_param = actual_param_owner.bias
                    logger.info(f"Key '{key}': Initialized bias for {type(actual_param_owner)}.")
                else:
                    skipped_keys_details.append(f"Key: {key}, Reason: Original bias is None and layer '{actual_param_owner}' is not Linear or LayerNorm. Cannot initialize.")
                    unpatched_keys_count +=1
                    continue
            
            # Special handling for RMSNorm which typically has no bias
            if isinstance(actual_param_owner, torch.nn.RMSNorm) and param_name_to_patch == "bias":
                skipped_keys_details.append(f"Key: {key}, Reason: Layer '{actual_param_owner}' is RMSNorm which has no bias parameter. Skipping bias diff.")
                unpatched_keys_count +=1
                continue


            if original_param is not None:
                if original_param.shape != diff_tensor.shape:
                    skipped_keys_details.append(f"Key: {key}, Reason: Shape mismatch. Model param: {original_param.shape}, LoRA diff: {diff_tensor.shape}. Layer: {actual_param_owner}")
                    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}'") # Too verbose, will log summary
                patched_keys_count += 1
            else:
                skipped_keys_details.append(f"Key: {key}, Reason: Original parameter '{param_name_to_patch}' is None and was not initialized. Layer: {actual_param_owner}")
                unpatched_keys_count += 1

        except AttributeError as e:
            skipped_keys_details.append(f"Key: {key}, Reason: AttributeError - {e}")
            unpatched_keys_count += 1
        except Exception as e:
            skipped_keys_details.append(f"Key: {key}, Reason: General Exception - {e}")
            unpatched_keys_count += 1

    logger.info(f"Manual patching summary: {patched_keys_count} keys patched, {unpatched_keys_count} keys failed or skipped.")
    if unpatched_keys_count > 0:
        logger.warning("Details of unpatched/skipped keys:")
        for detail in skipped_keys_details:
            logger.warning(f"  - {detail}")

# --- 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[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name
        
    export_to_video(output_frames_list, 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,
            steps,
            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)