Spaces:
Sleeping
Sleeping
| import copy | |
| import torch | |
| _orig_torch_load = torch.load | |
| def _torch_load(*args, **kwargs): | |
| kwargs.setdefault('weights_only', False) | |
| return _orig_torch_load(*args, **kwargs) | |
| torch.load = _torch_load | |
| from safetensors.torch import load_file | |
| from diffusers import StableDiffusionPipeline, DDIMScheduler, UNet2DConditionModel | |
| def get_module_kohya_state_dict(module, prefix: str, dtype: torch.dtype, adapter_name: str = "default"): | |
| kohya_ss_state_dict = {} | |
| for peft_key, weight in module.items(): | |
| kohya_key = peft_key.replace("unet.base_model.model", prefix) | |
| kohya_key = kohya_key.replace("lora_A", "lora_down") | |
| kohya_key = kohya_key.replace("lora_B", "lora_up") | |
| kohya_key = kohya_key.replace(".", "_", kohya_key.count(".") - 2) | |
| kohya_ss_state_dict[kohya_key] = weight.to(dtype) | |
| # Set alpha parameter | |
| if "lora_down" in kohya_key: | |
| alpha_key = f'{kohya_key.split(".")[0]}.alpha' | |
| kohya_ss_state_dict[alpha_key] = torch.tensor(8).to(dtype) | |
| return kohya_ss_state_dict | |
| def _make_pipe(vae, text_encoder, tokenizer, unet, scheduler): | |
| """Build a StableDiffusionPipeline around one specific UNet. | |
| The VAE / text encoder / tokenizer are shared (read-only at inference | |
| time), but every pipeline gets its OWN UNet module so the reverse and | |
| forward LoRA deltas are fused into separate weight sets — the same end | |
| state as the original copy.deepcopy(pipe) implementation. | |
| Everything stays on CPU while building: ZeroGPU rejects copy.deepcopy of | |
| pipelines whose tensors sit on a mix of devices, so devices are assigned | |
| only once, at the very end of loading (see load_consistency). | |
| """ | |
| return StableDiffusionPipeline( | |
| vae=vae, | |
| text_encoder=text_encoder, | |
| tokenizer=tokenizer, | |
| unet=unet, | |
| scheduler=scheduler, | |
| safety_checker=None, | |
| feature_extractor=None, | |
| requires_safety_checker=False, | |
| ) | |
| # Load models (DM, CM) | |
| def load_models(model_id, | |
| device, | |
| reverse_checkpoint, | |
| forward_checkpoint, | |
| r=64, | |
| w_embed_dim=0, | |
| teacher_checkpoint=None, | |
| dtype='fp32', | |
| ): | |
| # Diffusion | |
| # ------------------------------------------------------------ | |
| dtype = torch.float32 if dtype == 'fp32' else torch.float16 | |
| scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, | |
| set_alpha_to_one=False) | |
| ldm_stable = StableDiffusionPipeline.from_pretrained( | |
| model_id, scheduler=scheduler, safety_checker=None, | |
| requires_safety_checker=False).to('cpu', dtype=dtype) | |
| ldm_stable.enable_attention_slicing() | |
| # ------------------------------------------------------------ | |
| # Reverse consistency | |
| # ------------------------------------------------------------ | |
| if w_embed_dim > 0: | |
| print(f'Forward CD is initialized with guidance embedding, dim {w_embed_dim}') | |
| unet = UNet2DConditionModel.from_pretrained( | |
| model_id, subfolder="unet", | |
| time_cond_proj_dim=w_embed_dim, low_cpu_mem_usage=False, device_map=None, | |
| attn_implementation="sdpa" | |
| ).to('cpu') | |
| if teacher_checkpoint is not None: | |
| print(f'Embedded model is loading from {teacher_checkpoint}') | |
| unet.load_state_dict(torch.load(teacher_checkpoint, map_location='cpu')) | |
| ldm_stable.unet = unet | |
| ldm_stable.to(dtype=dtype) | |
| else: | |
| print('PROVIDE TEACHER') | |
| else: | |
| unet = UNet2DConditionModel.from_pretrained( | |
| model_id, subfolder="unet", attn_implementation="sdpa" | |
| ).to('cpu') | |
| if reverse_checkpoint is not None: | |
| print(f'Reverse CD is loading from {reverse_checkpoint}') | |
| reverse_cons_model = _make_pipe( | |
| ldm_stable.vae, ldm_stable.text_encoder, ldm_stable.tokenizer, | |
| copy.deepcopy(unet), ldm_stable.scheduler) | |
| lora_weight = load_file(reverse_checkpoint) | |
| lora_state_dict = get_module_kohya_state_dict(lora_weight, "lora_unet", torch.float16) | |
| reverse_cons_model.load_lora_weights(lora_state_dict) | |
| reverse_cons_model.fuse_lora() | |
| reverse_cons_model.to(dtype=dtype) | |
| else: | |
| reverse_cons_model = None | |
| # ------------------------------------------------------------ | |
| # Inverse consistency | |
| # ------------------------------------------------------------ | |
| if forward_checkpoint is not None: | |
| print(f'Forward CD is loading from {forward_checkpoint}') | |
| forward_cons_model = _make_pipe( | |
| ldm_stable.vae, ldm_stable.text_encoder, ldm_stable.tokenizer, | |
| copy.deepcopy(unet), ldm_stable.scheduler) | |
| lora_weight = load_file(forward_checkpoint) | |
| lora_state_dict = get_module_kohya_state_dict(lora_weight, "lora_unet", torch.float16) | |
| forward_cons_model.load_lora_weights(lora_state_dict) | |
| forward_cons_model.fuse_lora() | |
| forward_cons_model.to(dtype=dtype) | |
| else: | |
| forward_cons_model = None | |
| # ------------------------------------------------------------ | |
| return ldm_stable, reverse_cons_model, forward_cons_model | |