Update app.py

app.py

@@ -192,7 +192,7 @@ def _custom_convert_non_diffusers_wan_lora_to_diffusers(state_dict):
     return final_peft_state_dict
 
 
-def apply_manual_diff_patches(pipe_model, patches):
+def apply_manual_diff_patches(pipe_model: torch.nn.Module, patches: Dict[str, torch.Tensor]):
     """
     Manually applies diff_b/diff patches to the model.
     Assumes PEFT LoRA layers have already been loaded.
@@ -204,87 +204,95 @@ def apply_manual_diff_patches(pipe_model, patches):
     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:
-            module_to_patch = pipe_model
-            attrs = key.split(".")
+            # 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
-            # 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)
+            # 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:
-                    # 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}")
+                    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
             
-            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(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
 
-            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.")
+            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
 
-            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
+                    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}'")
-                patched_keys_count +=1
+                # logger.info(f"Successfully applied diff to '{key}'") # Too verbose, will log summary
+                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
-
+                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:
-            logger.error(f"AttributeError: Could not find module or parameter for key '{key}'. Error: {e}. Skipping.")
-            unpatched_keys_count +=1
+            skipped_keys_details.append(f"Key: {key}, Reason: AttributeError - {e}")
+            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.")
+            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}...")
@@ -411,6 +419,7 @@ with gr.Blocks() as demo:
             width_input,
             num_frames_input,
             guidance_scale_input,
+            steps,
             fps_input
         ],
         outputs=video_output