import torch import torch.nn as nn from typing import Dict, Tuple, Optional, Union import warnings try: from safetensors.torch import save_file as save_safetensors SAFETENSORS_AVAILABLE = True except ImportError: SAFETENSORS_AVAILABLE = False warnings.warn("safetensors not available. Install with: pip install safetensors") class LoRAExtractor: """ Extract LoRA tensors from the difference between original and fine-tuned models. LoRA (Low-Rank Adaptation) decomposes weight updates as ΔW = B @ A where: - A (lora_down): [rank, input_dim] matrix (saved as diffusion_model.param_name.lora_down.weight) - B (lora_up): [output_dim, rank] matrix (saved as diffusion_model.param_name.lora_up.weight) The decomposition uses SVD: ΔW = U @ S @ V^T ≈ (U @ S) @ V^T where: - lora_up = U @ S (contains all singular values) - lora_down = V^T (orthogonal matrix) Parameter handling based on name AND dimension: - 2D weight tensors: LoRA decomposition (.lora_down.weight, .lora_up.weight) - Any bias tensors: direct difference (.diff_b) - Other weight tensors (1D, 3D, 4D): full difference (.diff) Progress tracking and test mode are available for format validation and debugging. """ def __init__(self, rank: int = 128, threshold: float = 1e-6, test_mode: bool = False, show_reconstruction_errors: bool = False): """ Initialize LoRA extractor. Args: rank: Target rank for LoRA decomposition (default: 128) threshold: Minimum singular value threshold for decomposition test_mode: If True, creates zero tensors without computation for format testing show_reconstruction_errors: If True, calculates and displays reconstruction error for each LoRA pair """ self.rank = rank self.threshold = threshold self.test_mode = test_mode self.show_reconstruction_errors = show_reconstruction_errors def extract_lora_from_state_dicts( self, original_state_dict: Dict[str, torch.Tensor], finetuned_state_dict: Dict[str, torch.Tensor], device: str = 'cpu', show_progress: bool = True ) -> Dict[str, torch.Tensor]: """ Extract LoRA tensors for all matching parameters between two state dictionaries. Args: original_state_dict: State dict of the original model finetuned_state_dict: State dict of the fine-tuned model device: Device to perform computations on show_progress: Whether to display progress information Returns: Dictionary mapping parameter names to their LoRA components: - For 2D weight tensors: 'diffusion_model.layer.lora_down.weight', 'diffusion_model.layer.lora_up.weight' - For any bias tensors: 'diffusion_model.layer.diff_b' - For other weight tensors (1D, 3D, 4D): 'diffusion_model.layer.diff' """ lora_tensors = {} # Find common parameters and sort alphabetically for consistent processing order common_keys = sorted(set(original_state_dict.keys()) & set(finetuned_state_dict.keys())) total_params = len(common_keys) processed_params = 0 extracted_components = 0 if show_progress: print(f"Starting LoRA extraction for {total_params} parameters on {device}...") # Pre-move threshold to device for faster comparisons threshold_tensor = torch.tensor(self.threshold, device=device) for param_name in common_keys: if show_progress: processed_params += 1 progress_pct = (processed_params / total_params) * 100 print(f"[{processed_params:4d}/{total_params}] ({progress_pct:5.1f}%) Processing: {param_name}") # Move tensors to device once original_tensor = original_state_dict[param_name] finetuned_tensor = finetuned_state_dict[param_name] # Check if tensors have the same shape before moving to device if original_tensor.shape != finetuned_tensor.shape: if show_progress: print(f" → Shape mismatch: {original_tensor.shape} vs {finetuned_tensor.shape}. Skipping.") continue # Move to device and compute difference in one go for efficiency (skip in test mode) if not self.test_mode: if original_tensor.device != torch.device(device): original_tensor = original_tensor.to(device, non_blocking=True) if finetuned_tensor.device != torch.device(device): finetuned_tensor = finetuned_tensor.to(device, non_blocking=True) # Compute difference on device delta_tensor = finetuned_tensor - original_tensor # Fast GPU-based threshold check max_abs_diff = torch.max(torch.abs(delta_tensor)) if max_abs_diff <= threshold_tensor: if show_progress: print(f" → No significant changes detected (max diff: {max_abs_diff:.2e}), skipping") continue else: # Test mode - create dummy delta tensor with original shape and dtype delta_tensor = torch.zeros_like(original_tensor) if device != 'cpu': delta_tensor = delta_tensor.to(device) # Extract LoRA components based on tensor dimensionality extracted_tensors = self._extract_lora_components(delta_tensor, param_name) if extracted_tensors: lora_tensors.update(extracted_tensors) extracted_components += len(extracted_tensors) if show_progress: # Show meaningful component names instead of just 'weight' component_names = [] for key in extracted_tensors.keys(): if key.endswith('.lora_down.weight'): component_names.append('lora_down') elif key.endswith('.lora_up.weight'): component_names.append('lora_up') elif key.endswith('.diff_b'): component_names.append('diff_b') elif key.endswith('.diff'): component_names.append('diff') else: component_names.append(key.split('.')[-1]) print(f" → Extracted {len(extracted_tensors)} components: {component_names}") if show_progress: print(f"\nExtraction completed!") print(f"Processed: {processed_params}/{total_params} parameters") print(f"Extracted: {extracted_components} LoRA components") print(f"LoRA rank: {self.rank}") # Summary by type lora_down_count = sum(1 for k in lora_tensors.keys() if k.endswith('.lora_down.weight')) lora_up_count = sum(1 for k in lora_tensors.keys() if k.endswith('.lora_up.weight')) diff_b_count = sum(1 for k in lora_tensors.keys() if k.endswith('.diff_b')) diff_count = sum(1 for k in lora_tensors.keys() if k.endswith('.diff')) print(f"Summary: {lora_down_count} lora_down, {lora_up_count} lora_up, {diff_b_count} diff_b, {diff_count} diff") return lora_tensors def _extract_lora_components( self, delta_tensor: torch.Tensor, param_name: str ) -> Optional[Dict[str, torch.Tensor]]: """ Extract LoRA components from a delta tensor. Args: delta_tensor: Difference between fine-tuned and original tensor param_name: Name of the parameter (for generating output keys) Returns: Dictionary with modified parameter names as keys and tensors as values """ # Determine if this is a weight or bias parameter from the original name is_weight = 'weight' in param_name.lower() is_bias = 'bias' in param_name.lower() # Remove .weight or .bias suffix from parameter name base_name = param_name if base_name.endswith('.weight'): base_name = base_name[:-7] # Remove '.weight' elif base_name.endswith('.bias'): base_name = base_name[:-5] # Remove '.bias' # Add diffusion_model prefix base_name = f"diffusion_model.{base_name}" if self.test_mode: # Fast test mode - create zero tensors without computation if delta_tensor.dim() == 2 and is_weight: # 2D weight tensor -> LoRA decomposition output_dim, input_dim = delta_tensor.shape rank = min(self.rank, min(input_dim, output_dim)) return { f"{base_name}.lora_down.weight": torch.zeros(rank, input_dim, dtype=delta_tensor.dtype, device=delta_tensor.device), f"{base_name}.lora_up.weight": torch.zeros(output_dim, rank, dtype=delta_tensor.dtype, device=delta_tensor.device) } elif is_bias: # Any bias tensor (1D, 2D, etc.) -> .diff_b return {f"{base_name}.diff_b": torch.zeros_like(delta_tensor)} else: # Any weight tensor that's not 2D, or other tensors -> .diff return {f"{base_name}.diff": torch.zeros_like(delta_tensor)} # Normal mode - check dimensions AND parameter type if delta_tensor.dim() == 2 and is_weight: # 2D weight tensor (linear layer weight) - apply SVD decomposition return self._decompose_2d_tensor(delta_tensor, base_name) elif is_bias: # Any bias tensor (regardless of dimension) - save as .diff_b return {f"{base_name}.diff_b": delta_tensor.clone()} else: # Any other tensor (weight tensors that are 1D, 3D, 4D, or unknown tensors) - save as .diff return {f"{base_name}.diff": delta_tensor.clone()} def _decompose_2d_tensor(self, delta_tensor: torch.Tensor, base_name: str) -> Dict[str, torch.Tensor]: """ Decompose a 2D tensor using SVD on GPU for maximum performance. Args: delta_tensor: 2D tensor to decompose (output_dim × input_dim) base_name: Base name for the parameter (already processed, with diffusion_model prefix) Returns: Dictionary with lora_down and lora_up tensors: - lora_down: [rank, input_dim] - lora_up: [output_dim, rank] """ # Store original dtype and device dtype = delta_tensor.dtype device = delta_tensor.device # Perform SVD in float32 for numerical stability, but keep on same device delta_float = delta_tensor.float() if delta_tensor.dtype != torch.float32 else delta_tensor U, S, Vt = torch.linalg.svd(delta_float, full_matrices=False) # Determine effective rank (number of significant singular values) # Use GPU-accelerated operations significant_mask = S > self.threshold effective_rank = min(self.rank, torch.sum(significant_mask).item()) effective_rank = self.rank if effective_rank == 0: warnings.warn(f"No significant singular values found for {base_name}") effective_rank = 1 # Create LoRA matrices with correct SVD decomposition # Standard approach: put all singular values in lora_up, leave lora_down as V^T # This ensures: lora_up @ lora_down = (U @ S) @ V^T = U @ S @ V^T = ΔW ✓ lora_up = U[:, :effective_rank] * S[:effective_rank].unsqueeze(0) # [output_dim, rank] lora_down = Vt[:effective_rank, :] # [rank, input_dim] # Convert back to original dtype (keeping on same device) lora_up = lora_up.to(dtype) lora_down = lora_down.to(dtype) # Calculate and display reconstruction error if requested if self.show_reconstruction_errors: with torch.no_grad(): # Reconstruct the original delta tensor reconstructed = lora_up @ lora_down # Calculate various error metrics mse_error = torch.mean((delta_tensor - reconstructed) ** 2).item() max_error = torch.max(torch.abs(delta_tensor - reconstructed)).item() # Relative error original_norm = torch.norm(delta_tensor).item() relative_error = (torch.norm(delta_tensor - reconstructed).item() / original_norm * 100) if original_norm > 0 else 0 # Cosine similarity delta_flat = delta_tensor.flatten() reconstructed_flat = reconstructed.flatten() if torch.norm(delta_flat) > 0 and torch.norm(reconstructed_flat) > 0: cosine_sim = torch.nn.functional.cosine_similarity( delta_flat.unsqueeze(0), reconstructed_flat.unsqueeze(0) ).item() else: cosine_sim = 0.0 # Extract parameter name for display (remove diffusion_model prefix) display_name = base_name[16:] if base_name.startswith('diffusion_model.') else base_name print(f" LoRA Error [{display_name}]: MSE={mse_error:.2e}, Max={max_error:.2e}, Rel={relative_error:.2f}%, Cos={cosine_sim:.4f}, Rank={effective_rank}") return { f"{base_name}.lora_down.weight": lora_down, f"{base_name}.lora_up.weight": lora_up } def verify_reconstruction( self, lora_tensors: Dict[str, torch.Tensor], original_deltas: Dict[str, torch.Tensor] ) -> Dict[str, float]: """ Verify the quality of LoRA reconstruction for 2D tensors. Args: lora_tensors: Dictionary with LoRA tensors (flat structure with diffusion_model prefix) original_deltas: Dictionary with original delta tensors (without prefix) Returns: Dictionary mapping parameter names to reconstruction errors """ reconstruction_errors = {} # Group LoRA components by base parameter name lora_pairs = {} for key, tensor in lora_tensors.items(): if key.endswith('.lora_down.weight'): base_name = key[:-18] # Remove '.lora_down.weight' # Remove diffusion_model prefix for matching with original_deltas if base_name.startswith('diffusion_model.'): original_key = base_name[16:] # Remove 'diffusion_model.' else: original_key = base_name if base_name not in lora_pairs: lora_pairs[base_name] = {'original_key': original_key} lora_pairs[base_name]['lora_down'] = tensor elif key.endswith('.lora_up.weight'): base_name = key[:-16] # Remove '.lora_up.weight' # Remove diffusion_model prefix for matching with original_deltas if base_name.startswith('diffusion_model.'): original_key = base_name[16:] # Remove 'diffusion_model.' else: original_key = base_name if base_name not in lora_pairs: lora_pairs[base_name] = {'original_key': original_key} lora_pairs[base_name]['lora_up'] = tensor # Verify reconstruction for each complete LoRA pair for base_name, components in lora_pairs.items(): if 'lora_down' in components and 'lora_up' in components and 'original_key' in components: original_key = components['original_key'] if original_key in original_deltas: lora_down = components['lora_down'] lora_up = components['lora_up'] original_delta = original_deltas[original_key] # Get effective rank from the actual tensor dimensions effective_rank = min(lora_up.shape[1], lora_down.shape[0]) # Reconstruct: ΔW = lora_up @ lora_down (no additional scaling needed since it's built into lora_up) reconstructed = lora_up @ lora_down # Compute reconstruction error mse_error = torch.mean((original_delta - reconstructed) ** 2).item() reconstruction_errors[base_name] = mse_error return reconstruction_errors def compute_reconstruction_errors( original_tensor: torch.Tensor, reconstructed_tensor: torch.Tensor, target_tensor: torch.Tensor ) -> Dict[str, float]: """ Compute various error metrics between original, reconstructed, and target tensors. Args: original_tensor: Original tensor before fine-tuning reconstructed_tensor: Reconstructed tensor from LoRA (original + LoRA_reconstruction) target_tensor: Target tensor (fine-tuned) Returns: Dictionary with error metrics """ # Ensure all tensors are on the same device and have the same shape device = original_tensor.device reconstructed_tensor = reconstructed_tensor.to(device) target_tensor = target_tensor.to(device) # Compute differences delta_original = target_tensor - original_tensor # True fine-tuning difference delta_reconstructed = reconstructed_tensor - original_tensor # LoRA reconstructed difference reconstruction_error = target_tensor - reconstructed_tensor # Final reconstruction error # Compute various error metrics errors = {} # Mean Squared Error (MSE) errors['mse_delta'] = torch.mean((delta_original - delta_reconstructed) ** 2).item() errors['mse_final'] = torch.mean(reconstruction_error ** 2).item() # Mean Absolute Error (MAE) errors['mae_delta'] = torch.mean(torch.abs(delta_original - delta_reconstructed)).item() errors['mae_final'] = torch.mean(torch.abs(reconstruction_error)).item() # Relative errors (as percentages) original_norm = torch.norm(original_tensor).item() target_norm = torch.norm(target_tensor).item() delta_norm = torch.norm(delta_original).item() if original_norm > 0: errors['relative_error_original'] = (torch.norm(reconstruction_error).item() / original_norm) * 100 if target_norm > 0: errors['relative_error_target'] = (torch.norm(reconstruction_error).item() / target_norm) * 100 if delta_norm > 0: errors['relative_error_delta'] = (torch.norm(delta_original - delta_reconstructed).item() / delta_norm) * 100 # Cosine similarity (higher is better, 1.0 = perfect) delta_flat = delta_original.flatten() reconstructed_flat = delta_reconstructed.flatten() if torch.norm(delta_flat) > 0 and torch.norm(reconstructed_flat) > 0: cosine_sim = torch.nn.functional.cosine_similarity( delta_flat.unsqueeze(0), reconstructed_flat.unsqueeze(0) ).item() errors['cosine_similarity'] = cosine_sim else: errors['cosine_similarity'] = 0.0 # Signal-to-noise ratio (SNR) in dB if errors['mse_final'] > 0: signal_power = torch.mean(target_tensor ** 2).item() errors['snr_db'] = 10 * torch.log10(signal_power / errors['mse_final']).item() else: errors['snr_db'] = float('inf') return errors # Example usage and utility functions def load_and_extract_lora( original_model_path: str, finetuned_model_path: str, rank: int = 128, device: str = 'cuda' if torch.cuda.is_available() else 'cpu', show_progress: bool = True, test_mode: bool = False, show_reconstruction_errors: bool = False ) -> Dict[str, torch.Tensor]: """ Convenience function to load models and extract LoRA tensors with GPU acceleration. Args: original_model_path: Path to original model state dict finetuned_model_path: Path to fine-tuned model state dict rank: Target LoRA rank (default: 128) device: Device for computation (defaults to GPU if available) show_progress: Whether to display progress information test_mode: If True, creates zero tensors without computation for format testing show_reconstruction_errors: If True, calculates and displays reconstruction error for each LoRA pair Returns: Dictionary of LoRA tensors with modified parameter names as keys """ # Load state dictionaries directly to CPU first (safetensors loads to CPU by default) if show_progress: print(f"Loading original model from: {original_model_path}") original_state_dict = torch.load(original_model_path, map_location='cpu') if show_progress: print(f"Loading fine-tuned model from: {finetuned_model_path}") finetuned_state_dict = torch.load(finetuned_model_path, map_location='cpu') # Handle nested state dicts (if wrapped in 'model' key or similar) if 'state_dict' in original_state_dict: original_state_dict = original_state_dict['state_dict'] if 'state_dict' in finetuned_state_dict: finetuned_state_dict = finetuned_state_dict['state_dict'] # Extract LoRA tensors with GPU acceleration extractor = LoRAExtractor(rank=rank, test_mode=test_mode, show_reconstruction_errors=show_reconstruction_errors) lora_tensors = extractor.extract_lora_from_state_dicts( original_state_dict, finetuned_state_dict, device=device, show_progress=show_progress ) return lora_tensors def save_lora_tensors(lora_tensors: Dict[str, torch.Tensor], save_path: str): """Save extracted LoRA tensors to disk.""" torch.save(lora_tensors, save_path) print(f"LoRA tensors saved to {save_path}") def save_lora_safetensors(lora_tensors: Dict[str, torch.Tensor], save_path: str, rank: int = None): """Save extracted LoRA tensors as safetensors format with metadata.""" if not SAFETENSORS_AVAILABLE: raise ImportError("safetensors not available. Install with: pip install safetensors") # Ensure all tensors are contiguous for safetensors contiguous_tensors = {k: v.contiguous() if v.is_floating_point() else v.contiguous() for k, v in lora_tensors.items()} # Add rank as metadata if provided metadata = {} if rank is not None: metadata["rank"] = str(rank) save_safetensors(contiguous_tensors, save_path, metadata=metadata if metadata else None) print(f"LoRA tensors saved as safetensors to {save_path}") if metadata: print(f"Metadata: {metadata}") def analyze_lora_tensors(lora_tensors: Dict[str, torch.Tensor]): """Analyze the extracted LoRA tensors.""" print(f"Extracted LoRA tensors ({len(lora_tensors)} components):") # Group by type for better organization lora_down_tensors = {k: v for k, v in lora_tensors.items() if k.endswith('.lora_down.weight')} lora_up_tensors = {k: v for k, v in lora_tensors.items() if k.endswith('.lora_up.weight')} diff_b_tensors = {k: v for k, v in lora_tensors.items() if k.endswith('.diff_b')} diff_tensors = {k: v for k, v in lora_tensors.items() if k.endswith('.diff')} if lora_down_tensors: print(f"\nLinear LoRA down matrices ({len(lora_down_tensors)}):") for name, tensor in lora_down_tensors.items(): print(f" {name}: {tensor.shape}") if lora_up_tensors: print(f"\nLinear LoRA up matrices ({len(lora_up_tensors)}):") for name, tensor in lora_up_tensors.items(): print(f" {name}: {tensor.shape}") if diff_b_tensors: print(f"\nBias differences ({len(diff_b_tensors)}):") for name, tensor in diff_b_tensors.items(): print(f" {name}: {tensor.shape}") if diff_tensors: print(f"\nFull weight differences ({len(diff_tensors)}):") print(" (Includes conv, modulation, and other multi-dimensional tensors)") for name, tensor in diff_tensors.items(): print(f" {name}: {tensor.shape}") # Example usage if __name__ == "__main__": from safetensors.torch import load_file as load_safetensors # Load original and fine-tuned models from safetensors files original_state_dict = load_safetensors("ckpts/wan2.2_text2video_14B_high_mbf16.safetensors") finetuned_state_dict = load_safetensors("ckpts/wan2.2_text2video_14B_low_mbf16.safetensors") # original_state_dict = load_safetensors("ckpts/flux1-dev_bf16.safetensors") # finetuned_state_dict = load_safetensors("ckpts/flux1-schnell_bf16.safetensors") print(f"Loaded original model with {len(original_state_dict)} parameters") print(f"Loaded fine-tuned model with {len(finetuned_state_dict)} parameters") # extractor_test = LoRAExtractor(test_mode=True) extractor_test = LoRAExtractor(show_reconstruction_errors=True, rank=128) lora_tensors_test = extractor_test.extract_lora_from_state_dicts( original_state_dict, finetuned_state_dict, device='cuda', show_progress=True ) print("\nTest mode tensor keys (first 10):") for i, key in enumerate(sorted(lora_tensors_test.keys())): if i < 10: print(f" {key}: {lora_tensors_test[key].shape}") elif i == 10: print(f" ... and {len(lora_tensors_test) - 10} more") break # Always save as extracted_lora.safetensors for easier testing save_lora_safetensors(lora_tensors_test, "extracted_lora.safetensors")