Spaces:
Sleeping
Sleeping
File size: 5,311 Bytes
46aa7f4 dc7f251 cf95a0d dc7f251 46aa7f4 cf95a0d 46aa7f4 cf95a0d 46aa7f4 cf95a0d 46aa7f4 cf95a0d dc7f251 cf95a0d dc7f251 cf95a0d 46aa7f4 dc7f251 46aa7f4 cf95a0d dc7f251 cf95a0d 46aa7f4 dc7f251 46aa7f4 dc7f251 46aa7f4 dc7f251 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | 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
|