| import os |
| import glob |
| import torch |
| import torch.distributed as dist |
| from collections.abc import Sequence |
|
|
| from src.utils.fsdp_load import maybe_load_fsdp_model, pt_weights_iterator, safetensors_weights_iterator |
| from src.utils.logging import get_logger |
| from src.utils.lora import validate_lora_state_dict |
| from src.utils.constants import PRECISION_TO_TYPE |
| from src.utils.utils import build_from_config, get_obj_from_str |
|
|
|
|
| def _uses_pretrained_loading(config: dict | None) -> bool: |
| if config is None or config.get("pretrained") is None: |
| return False |
|
|
| target = str(config.get("target", "")) |
| try: |
| cls = get_obj_from_str(target) |
| except Exception: |
| return target.startswith("diffusers.") |
|
|
| return hasattr(cls, "from_pretrained") |
|
|
|
|
| def _strip_common_state_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: |
| for prefix in ("model.", "module.", "transformer."): |
| stripped_state_dict = {} |
| matched = 0 |
| for key, value in state_dict.items(): |
| if key.startswith(prefix): |
| stripped_state_dict[key[len(prefix):]] = value |
| matched += 1 |
| else: |
| stripped_state_dict[key] = value |
| if matched > 0: |
| state_dict = stripped_state_dict |
| return state_dict |
|
|
|
|
| def _normalize_train_modules(train_modules: str | Sequence[str] | None) -> list[str] | None: |
| if train_modules is None: |
| return None |
| if isinstance(train_modules, str): |
| train_modules = [train_modules] |
|
|
| normalized = [module_name.strip() for module_name in train_modules if module_name and module_name.strip()] |
| return normalized or None |
|
|
|
|
| def _configure_dit_trainable_modules(model: torch.nn.Module, cfg) -> None: |
| logger = get_logger() |
| train_modules = _normalize_train_modules(getattr(cfg, "dit_train_modules", None)) |
|
|
| if train_modules is None: |
| model.requires_grad_(True) |
| return |
|
|
| model.requires_grad_(False) |
| named_modules = dict(model.named_modules()) |
| missing_modules = [module_name for module_name in train_modules if module_name not in named_modules] |
| if missing_modules: |
| available_modules = [module_name for module_name in named_modules.keys() if module_name] |
| preview = ", ".join(available_modules[:20]) |
| if len(available_modules) > 20: |
| preview = f"{preview}, ..." |
| raise ValueError( |
| f"Unknown dit_train_modules: {missing_modules}. " |
| f"Available DiT modules include: {preview}" |
| ) |
|
|
| for module_name in train_modules: |
| named_modules[module_name].requires_grad_(True) |
|
|
| trainable_param_names = [name for name, param in model.named_parameters() if param.requires_grad] |
| if not trainable_param_names: |
| raise ValueError(f"No trainable DiT parameters were enabled for dit_train_modules={train_modules}.") |
|
|
| logger.info(f"Training {len(trainable_param_names)} parameter tensors.") |
|
|
|
|
| def _load_vae(cfg, device: torch.device): |
| """Load VAE weights and initialize complete modules absent from the checkpoint.""" |
| vae_kwargs = {"torch_dtype": PRECISION_TO_TYPE[cfg.vae_precision], "device": device} |
| if not _uses_pretrained_loading(cfg.vae_arch_config): |
| return build_from_config(cfg.vae_arch_config, **vae_kwargs) |
| vae, loading = build_from_config( |
| cfg.vae_arch_config, **vae_kwargs, low_cpu_mem_usage=False, output_loading_info=True, |
| ) |
| missing_by_module = {} |
| for key in loading["missing_keys"]: |
| module_name, _, parameter_name = key.rpartition(".") |
| missing_by_module.setdefault(module_name, set()).add(parameter_name) |
| for module_name, missing in sorted(missing_by_module.items()): |
| module = vae.get_submodule(module_name) |
| parameters = set(dict(module.named_parameters(recurse=False))) |
| reset = getattr(module, "reset_parameters", None) |
| if missing != parameters or not callable(reset): |
| raise ValueError(f"Cannot initialize incomplete VAE module {module_name!r}: missing={sorted(missing)}.") |
| reset() |
| return vae |
|
|
|
|
| def load_pipeline(cfg, dit, device: torch.device): |
| """Build the frozen VAE, text encoder, scheduler, and editing pipeline.""" |
| vae = _load_vae(cfg, device).eval().requires_grad_(False) |
| tokenizer, text_encoder = build_from_config( |
| cfg.text_encoder_arch_config, torch_dtype=PRECISION_TO_TYPE[cfg.text_encoder_precision], device=device, |
| ) |
| pipeline = build_from_config( |
| cfg.pipeline_arch_config or {"target": "src.models.common.diffusion.pipelines.Pipeline"}, |
| vae=vae, |
| tokenizer=tokenizer, |
| text_encoder=text_encoder, |
| transformer=dit, |
| scheduler=build_from_config(cfg.scheduler_arch_config), |
| args=cfg, |
| ) |
| return pipeline.to(device) |
|
|
|
|
| def _read_dit_checkpoint(path, checkpoint_type): |
| if checkpoint_type == "pt": |
| state = dict(pt_weights_iterator([path])) |
| elif checkpoint_type == "safetensor": |
| files = [str(path)] if os.path.isfile(path) else sorted(glob.glob(os.path.join(str(path), "*.safetensors"))) |
| if not files: |
| raise FileNotFoundError(f"No safetensors checkpoints found at {path}.") |
| state = dict(safetensors_weights_iterator(files)) |
| else: |
| raise ValueError(f"Unknown checkpoint type: {checkpoint_type!r}.") |
| return state.get("model", state) |
|
|
|
|
| def _merge_generator_lora(model, checkpoint, cfg): |
| from peft import LoraConfig |
| from peft.utils import set_peft_model_state_dict |
|
|
| weights = checkpoint["generator_lora"] |
| adapter_config = checkpoint.get("generator_lora_config") |
| if adapter_config is None: |
| ranks = {weight.shape[0] for name, weight in weights.items() if ".lora_A." in name} |
| if len(ranks) != 1: |
| raise ValueError("The generator checkpoint must contain a uniform-rank LoRA adapter.") |
| rank = ranks.pop() |
| adapter_config = { |
| "r": rank, |
| "lora_alpha": getattr(cfg, "dmd_lora_alpha", rank), |
| "target_modules": sorted({name.split(".lora_")[0] for name in weights if ".lora_" in name}), |
| "bias": getattr(cfg, "dmd_lora_bias", "none"), |
| } |
| adapter_name = checkpoint.get("generator_adapter_name", "generator") |
| model.add_adapter(LoraConfig(**adapter_config), adapter_name=adapter_name) |
| validate_lora_state_dict(model, weights, adapter_name) |
| result = set_peft_model_state_dict(model, weights, adapter_name=adapter_name) |
| missing_adapter_keys = [name for name in result.missing_keys if ".lora_" in name] |
| if missing_adapter_keys or result.unexpected_keys: |
| raise ValueError("Generator LoRA weights do not match the configured model.") |
| model.fuse_lora(adapter_names=[adapter_name], safe_fusing=True) |
| model.unload_lora() |
|
|
|
|
| def load_dit(cfg, device: torch.device, *, use_fsdp: bool = True) -> torch.nn.Module: |
| """Load DiT model with FSDP support.""" |
| logger = get_logger() |
|
|
| state_dict = _read_dit_checkpoint(cfg.dit_ckpt, cfg.dit_ckpt_type) if cfg.dit_ckpt is not None else None |
| adapter_checkpoint = None |
| if state_dict is not None and "generator_lora" in state_dict: |
| adapter_checkpoint = state_dict |
| base_path = getattr(cfg, "dmd_base_ckpt", None) or getattr(cfg, "dmd_stage1_ckpt", None) |
| if not base_path or os.path.abspath(base_path) == os.path.abspath(cfg.dit_ckpt): |
| raise ValueError("Set JOYAI_STAGE1_CKPT or dmd_base_ckpt to the SFT base for generator LoRA inference.") |
| state_dict = _read_dit_checkpoint(base_path, getattr(cfg, "dmd_base_ckpt_type", "pt")) |
|
|
| dtype = PRECISION_TO_TYPE[cfg.dit_precision] |
| use_pretrained_loading = _uses_pretrained_loading(cfg.dit_arch_config) |
| if use_pretrained_loading: |
| model_kwargs = {'torch_dtype': dtype, 'device': device} |
| else: |
| model_kwargs = {'dtype': dtype, 'device': device, 'args': cfg} |
| model = build_from_config(cfg.dit_arch_config, **model_kwargs) |
| if not use_fsdp or not dist.is_initialized() or dist.get_world_size() == 1: |
| model.to(device=device) |
|
|
| if state_dict is not None: |
| state_dict = _strip_common_state_prefixes(state_dict) |
|
|
| load_state_dict = {} |
| for k, v in state_dict.items(): |
| if ( |
| k == "img_in.weight" |
| and hasattr(model, "img_in") |
| and model.img_in.weight.shape != v.shape |
| ): |
| v = v.reshape(model.img_in.weight.shape) |
| load_state_dict[k] = v |
| strict = not use_pretrained_loading |
| missing_keys, unexpected_keys = model.load_state_dict(load_state_dict, strict=strict) |
| if missing_keys: |
| logger.warning(f"Missing keys when loading DiT: {missing_keys[:20]}") |
| if unexpected_keys: |
| logger.warning(f"Unexpected keys when loading DiT: {unexpected_keys[:20]}") |
|
|
| if adapter_checkpoint is not None: |
| _merge_generator_lora(model, adapter_checkpoint, cfg) |
| _configure_dit_trainable_modules(model, cfg) |
|
|
| if use_fsdp: |
| model = maybe_load_fsdp_model( |
| model=model, |
| hsdp_shard_dim=cfg.hsdp_shard_dim, |
| reshard_after_forward=cfg.reshard_after_forward, |
| param_dtype=dtype, |
| reduce_dtype=torch.float32, |
| output_dtype=None, |
| cpu_offload=cfg.cpu_offload, |
| pin_cpu_memory=cfg.pin_cpu_memory, |
| ) |
|
|
| total_params = sum(p.numel() for p in model.parameters()) |
| logger.info(f"Instantiate model with {total_params / 1e9:.2f}B parameters") |
|
|
| param_dtypes = {param.dtype for param in model.parameters()} |
| if len(param_dtypes) > 1: |
| logger.warning( |
| f"Model has mixed dtypes: {param_dtypes}. Converting to {dtype}") |
| model = model.to(dtype) |
|
|
| return model.eval() |
|
|
|
|
| __all__ = [ |
| "load_dit", |
| "load_pipeline", |
| ] |
|
|
