Spaces:
Running
on
T4
Running
on
T4
| from timeit import default_timer as timer | |
| from datetime import timedelta | |
| from PIL import Image | |
| import os | |
| import numpy as np | |
| from einops import rearrange | |
| import torch | |
| import torch.nn.functional as F | |
| from torchvision import transforms | |
| import transformers | |
| from accelerate import Accelerator | |
| from accelerate.utils import set_seed | |
| from packaging import version | |
| from PIL import Image | |
| import tqdm | |
| from transformers import AutoTokenizer, PretrainedConfig | |
| import diffusers | |
| from diffusers import ( | |
| AutoencoderKL, | |
| DDPMScheduler, | |
| DiffusionPipeline, | |
| DPMSolverMultistepScheduler, | |
| StableDiffusionPipeline, | |
| UNet2DConditionModel, | |
| ) | |
| from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin | |
| from diffusers.models.attention_processor import ( | |
| AttnAddedKVProcessor, | |
| AttnAddedKVProcessor2_0, | |
| LoRAAttnAddedKVProcessor, | |
| LoRAAttnProcessor, | |
| LoRAAttnProcessor2_0, | |
| SlicedAttnAddedKVProcessor, | |
| ) | |
| from diffusers.optimization import get_scheduler | |
| from diffusers.utils import check_min_version | |
| from diffusers.utils.import_utils import is_xformers_available | |
| # Will error if the minimal version of diffusers is not installed. Remove at your own risks. | |
| check_min_version("0.17.0") | |
| def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str, revision: str): | |
| text_encoder_config = PretrainedConfig.from_pretrained( | |
| pretrained_model_name_or_path, | |
| subfolder="text_encoder", | |
| revision=revision, | |
| ) | |
| model_class = text_encoder_config.architectures[0] | |
| if model_class == "CLIPTextModel": | |
| from transformers import CLIPTextModel | |
| return CLIPTextModel | |
| elif model_class == "RobertaSeriesModelWithTransformation": | |
| from diffusers.pipelines.alt_diffusion.modeling_roberta_series import RobertaSeriesModelWithTransformation | |
| return RobertaSeriesModelWithTransformation | |
| elif model_class == "T5EncoderModel": | |
| from transformers import T5EncoderModel | |
| return T5EncoderModel | |
| else: | |
| raise ValueError(f"{model_class} is not supported.") | |
| def tokenize_prompt(tokenizer, prompt, tokenizer_max_length=None): | |
| if tokenizer_max_length is not None: | |
| max_length = tokenizer_max_length | |
| else: | |
| max_length = tokenizer.model_max_length | |
| text_inputs = tokenizer( | |
| prompt, | |
| truncation=True, | |
| padding="max_length", | |
| max_length=max_length, | |
| return_tensors="pt", | |
| ) | |
| return text_inputs | |
| def encode_prompt(text_encoder, input_ids, attention_mask, text_encoder_use_attention_mask=False): | |
| text_input_ids = input_ids.to(text_encoder.device) | |
| if text_encoder_use_attention_mask: | |
| attention_mask = attention_mask.to(text_encoder.device) | |
| else: | |
| attention_mask = None | |
| prompt_embeds = text_encoder( | |
| text_input_ids, | |
| attention_mask=attention_mask, | |
| ) | |
| prompt_embeds = prompt_embeds[0] | |
| return prompt_embeds | |
| # model_path: path of the model | |
| # image: input image, have not been pre-processed | |
| # save_lora_dir: the path to save the lora | |
| # prompt: the user input prompt | |
| # lora_steps: number of lora training step | |
| # lora_lr: learning rate of lora training | |
| # lora_rank: the rank of lora | |
| def train_lora(image, prompt, save_lora_dir, model_path=None, tokenizer=None, text_encoder=None, vae=None, unet=None, noise_scheduler=None, lora_steps=200, lora_lr=2e-4, lora_rank=16, weight_name=None, safe_serialization=False, progress=tqdm): | |
| # initialize accelerator | |
| accelerator = Accelerator( | |
| gradient_accumulation_steps=1, | |
| # mixed_precision='fp16' | |
| ) | |
| set_seed(0) | |
| # Load the tokenizer | |
| if tokenizer is None: | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| model_path, | |
| subfolder="tokenizer", | |
| revision=None, | |
| use_fast=False, | |
| ) | |
| # initialize the model | |
| if noise_scheduler is None: | |
| noise_scheduler = DDPMScheduler.from_pretrained(model_path, subfolder="scheduler") | |
| if text_encoder is None: | |
| text_encoder_cls = import_model_class_from_model_name_or_path(model_path, revision=None) | |
| text_encoder = text_encoder_cls.from_pretrained( | |
| model_path, subfolder="text_encoder", revision=None | |
| ) | |
| if vae is None: | |
| vae = AutoencoderKL.from_pretrained( | |
| model_path, subfolder="vae", revision=None | |
| ) | |
| if unet is None: | |
| unet = UNet2DConditionModel.from_pretrained( | |
| model_path, subfolder="unet", revision=None | |
| ) | |
| # set device and dtype | |
| device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") | |
| vae.requires_grad_(False) | |
| text_encoder.requires_grad_(False) | |
| unet.requires_grad_(False) | |
| unet.to(device) | |
| vae.to(device) | |
| text_encoder.to(device) | |
| # initialize UNet LoRA | |
| unet_lora_attn_procs = {} | |
| for name, attn_processor in unet.attn_processors.items(): | |
| cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim | |
| if name.startswith("mid_block"): | |
| hidden_size = unet.config.block_out_channels[-1] | |
| elif name.startswith("up_blocks"): | |
| block_id = int(name[len("up_blocks.")]) | |
| hidden_size = list(reversed(unet.config.block_out_channels))[block_id] | |
| elif name.startswith("down_blocks"): | |
| block_id = int(name[len("down_blocks.")]) | |
| hidden_size = unet.config.block_out_channels[block_id] | |
| else: | |
| raise NotImplementedError("name must start with up_blocks, mid_blocks, or down_blocks") | |
| if isinstance(attn_processor, (AttnAddedKVProcessor, SlicedAttnAddedKVProcessor, AttnAddedKVProcessor2_0)): | |
| lora_attn_processor_class = LoRAAttnAddedKVProcessor | |
| else: | |
| lora_attn_processor_class = ( | |
| LoRAAttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else LoRAAttnProcessor | |
| ) | |
| unet_lora_attn_procs[name] = lora_attn_processor_class( | |
| hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=lora_rank | |
| ) | |
| unet.set_attn_processor(unet_lora_attn_procs) | |
| unet_lora_layers = AttnProcsLayers(unet.attn_processors) | |
| # Optimizer creation | |
| params_to_optimize = (unet_lora_layers.parameters()) | |
| optimizer = torch.optim.AdamW( | |
| params_to_optimize, | |
| lr=lora_lr, | |
| betas=(0.9, 0.999), | |
| weight_decay=1e-2, | |
| eps=1e-08, | |
| ) | |
| lr_scheduler = get_scheduler( | |
| "constant", | |
| optimizer=optimizer, | |
| num_warmup_steps=0, | |
| num_training_steps=lora_steps, | |
| num_cycles=1, | |
| power=1.0, | |
| ) | |
| # prepare accelerator | |
| unet_lora_layers = accelerator.prepare_model(unet_lora_layers) | |
| optimizer = accelerator.prepare_optimizer(optimizer) | |
| lr_scheduler = accelerator.prepare_scheduler(lr_scheduler) | |
| # initialize text embeddings | |
| with torch.no_grad(): | |
| text_inputs = tokenize_prompt(tokenizer, prompt, tokenizer_max_length=None) | |
| text_embedding = encode_prompt( | |
| text_encoder, | |
| text_inputs.input_ids, | |
| text_inputs.attention_mask, | |
| text_encoder_use_attention_mask=False | |
| ) | |
| if type(image) == np.ndarray: | |
| image = Image.fromarray(image) | |
| # initialize latent distribution | |
| image_transforms = transforms.Compose( | |
| [ | |
| transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR), | |
| # transforms.RandomCrop(512), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.5], [0.5]), | |
| ] | |
| ) | |
| image = image_transforms(image).to(device) | |
| image = image.unsqueeze(dim=0) | |
| latents_dist = vae.encode(image).latent_dist | |
| for _ in progress.tqdm(range(lora_steps), desc="Training LoRA..."): | |
| unet.train() | |
| model_input = latents_dist.sample() * vae.config.scaling_factor | |
| # Sample noise that we'll add to the latents | |
| noise = torch.randn_like(model_input) | |
| bsz, channels, height, width = model_input.shape | |
| # Sample a random timestep for each image | |
| timesteps = torch.randint( | |
| 0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device | |
| ) | |
| timesteps = timesteps.long() | |
| # Add noise to the model input according to the noise magnitude at each timestep | |
| # (this is the forward diffusion process) | |
| noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps) | |
| # Predict the noise residual | |
| model_pred = unet(noisy_model_input, timesteps, text_embedding).sample | |
| # Get the target for loss depending on the prediction type | |
| if noise_scheduler.config.prediction_type == "epsilon": | |
| target = noise | |
| elif noise_scheduler.config.prediction_type == "v_prediction": | |
| target = noise_scheduler.get_velocity(model_input, noise, timesteps) | |
| else: | |
| raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") | |
| loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean") | |
| accelerator.backward(loss) | |
| optimizer.step() | |
| lr_scheduler.step() | |
| optimizer.zero_grad() | |
| # save the trained lora | |
| # unet = unet.to(torch.float32) | |
| # vae = vae.to(torch.float32) | |
| # text_encoder = text_encoder.to(torch.float32) | |
| # unwrap_model is used to remove all special modules added when doing distributed training | |
| # so here, there is no need to call unwrap_model | |
| # unet_lora_layers = accelerator.unwrap_model(unet_lora_layers) | |
| LoraLoaderMixin.save_lora_weights( | |
| save_directory=save_lora_dir, | |
| unet_lora_layers=unet_lora_layers, | |
| text_encoder_lora_layers=None, | |
| weight_name=weight_name, | |
| safe_serialization=safe_serialization | |
| ) | |
| def load_lora(unet, lora_0, lora_1, alpha): | |
| lora = {} | |
| for key in lora_0: | |
| lora[key] = (1 - alpha) * lora_0[key] + alpha * lora_1[key] | |
| unet.load_attn_procs(lora) | |
| return unet | |
| # import safetensors | |
| # unet = UNet2DConditionModel.from_pretrained( | |
| # "stabilityai/stable-diffusion-2-1-base", subfolder="unet", revision=None | |
| # ) | |
| # lora = safetensors.torch.load_file("../models/lora/majicmixRealistic_betterV2V25.safetensors", device="cuda") | |
| # unet = safetensors.torch.load_file("../stabilityai/stable-diffusion-1-5/v1-5-pruned-emaonly.safetensors", device="cuda") | |
| # with open("lora.txt", "w") as f: | |
| # for key in lora: | |
| # f.write(f"{key} {lora[key].shape}\n") | |
| # with open("unet.txt", "w") as f: | |
| # for key in unet: | |
| # f.write(f"{key} {unet[key].shape}\n") | |
| # unet.load_attn_procs(lora) | |
| # lora_path = "models/lora" | |
| # image_path_1 = "input/sculpture.jpg" | |
| # # image_path_0 = "input/realdog0.jpg" | |
| # prompt = "a photo of a sculpture" | |
| # train_lora(Image.open(image_path_1), prompt, lora_path, "stabilityai/stable-diffusion-1-5", weight_name="sculpture_v15.safetensors", safe_serialization=True) | |
| # train_lora(image_path_0, prompt, "stabilityai/stable-diffusion-2-1-base", lora_path, weight_name="realdog0.ckpt") | |
| # realdog1_lora = torch.load(os.path.join(lora_path, "realdog1.ckpt")) | |
| # realdog0_lora = torch.load(os.path.join(lora_path, "realdog0.ckpt")) | |
| # pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float32) | |
| # pipe.to("cuda") | |
| # for t in torch.linspace(0, 1, 10): | |
| # lora = {} | |
| # for key in realdog0_lora: | |
| # lora[key] = (1 - t) * realdog1_lora[key] + t * realdog0_lora[key] | |
| # pipe.unet.load_attn_procs(lora) | |
| # image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0] | |
| # image.save(f"test/lora_interp/{t}.jpg") |