| 12 | return RUNNING_FLAG |
| 13 | |
| 14 | class WanLoraWrapper: |
| 15 | def __init__(self, wan_model): |
| 16 | self.model = wan_model |
| 17 | self.lora_metadata = {} |
| 18 | # self.override_dict = {} # On CPU |
| 19 | |
| 20 | def load_lora(self, lora_path, lora_name=None): |
| 21 | if lora_name is None: |
| 22 | lora_name = os.path.basename(lora_path).split(".")[0] |
| 23 | |
| 24 | if lora_name in self.lora_metadata: |
| 25 | logger.info(f"LoRA {lora_name} already loaded, skipping...") |
| 26 | return lora_name |
| 27 | |
| 28 | self.lora_metadata[lora_name] = {"path": lora_path} |
| 29 | logger.info(f"Registered LoRA metadata for: {lora_name} from {lora_path}") |
| 30 | |
| 31 | return lora_name |
| 32 | |
| 33 | def _load_lora_file(self, file_path, param_dtype): |
| 34 | with safe_open(file_path, framework="pt") as f: |
| 35 | tensor_dict = {key: f.get_tensor(key).to(param_dtype) for key in f.keys()} |
| 36 | return tensor_dict |
| 37 | |
| 38 | def apply_lora(self, lora_name, alpha=1.0, param_dtype=torch.bfloat16, device='cpu'): |
| 39 | if lora_name not in self.lora_metadata: |
| 40 | logger.info(f"LoRA {lora_name} not found. Please load it first.") |
| 41 | |
| 42 | |
| 43 | |
| 44 | lora_weights = self._load_lora_file(self.lora_metadata[lora_name]["path"], param_dtype) |
| 45 | # weight_dict = self.model.original_weight_dict |
| 46 | self._apply_lora_weights(lora_weights, alpha, device) |
| 47 | # self.model._init_weights(weight_dict) |
| 48 | |
| 49 | logger.info(f"Applied LoRA: {lora_name} with alpha={alpha}") |
| 50 | return True |
| 51 | |
| 52 | def get_parameter_by_name(self, model, param_name): |
| 53 | parts = param_name.split('.') |
| 54 | current = model |
| 55 | for part in parts: |
| 56 | if part.isdigit(): |
| 57 | current = current[int(part)] |
| 58 | else: |
| 59 | current = getattr(current, part) |
| 60 | return current |
| 61 | |
| 62 | @torch.no_grad() |
| 63 | def _apply_lora_weights(self, lora_weights, alpha, device): |
| 64 | lora_pairs = {} |
| 65 | prefix = "diffusion_model." |
| 66 | |
| 67 | for key in lora_weights.keys(): |
| 68 | if key.endswith("lora_down.weight") and key.startswith(prefix): |
| 69 | base_name = key[len(prefix) :].replace("lora_down.weight", "weight") |
| 70 | b_key = key.replace("lora_down.weight", "lora_up.weight") |
| 71 | if b_key in lora_weights: |