AnchoredAI/llm-grounded-diffusion

1import torch2from transformers import CLIPTextModel, CLIPTokenizer3from diffusers import AutoencoderKL, DDIMScheduler, DDIMInverseScheduler, DPMSolverMultistepScheduler4from .unet_2d_condition import UNet2DConditionModel5from easydict import EasyDict6import numpy as np7# For compatibility8from utils.latents import get_unscaled_latents, get_scaled_latents, blend_latents9from utils import torch_device10 11def load_sd(key="runwayml/stable-diffusion-v1-5", use_fp16=False, load_inverse_scheduler=True):12    """13    Keys:14     key = "CompVis/stable-diffusion-v1-4"15     key = "runwayml/stable-diffusion-v1-5"16     key = "stabilityai/stable-diffusion-2-1-base"17     18    Unpack with:19    ```20    model_dict = load_sd(key=key, use_fp16=use_fp16)21    vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict.scheduler, model_dict.dtype22    ```23    24    use_fp16: fp16 might have degraded performance25    """26    27    # run final results in fp3228    if use_fp16:29        dtype = torch.float1630        revision = "fp16"31    else:32        dtype = torch.float33        revision = "main"34        35    vae = AutoencoderKL.from_pretrained(key, subfolder="vae", revision=revision, torch_dtype=dtype).to(torch_device)36    tokenizer = CLIPTokenizer.from_pretrained(key, subfolder="tokenizer", revision=revision, torch_dtype=dtype)37    text_encoder = CLIPTextModel.from_pretrained(key, subfolder="text_encoder", revision=revision, torch_dtype=dtype).to(torch_device)38    unet = UNet2DConditionModel.from_pretrained(key, subfolder="unet", revision=revision, torch_dtype=dtype).to(torch_device)39    dpm_scheduler = DPMSolverMultistepScheduler.from_pretrained(key, subfolder="scheduler", revision=revision, torch_dtype=dtype)40    scheduler = DDIMScheduler.from_pretrained(key, subfolder="scheduler", revision=revision, torch_dtype=dtype)41 42    model_dict = EasyDict(vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, unet=unet, scheduler=scheduler, dpm_scheduler=dpm_scheduler, dtype=dtype)43    44    if load_inverse_scheduler:45        inverse_scheduler = DDIMInverseScheduler.from_config(scheduler.config)46        model_dict.inverse_scheduler = inverse_scheduler47    48    return model_dict49 50def encode_prompts(tokenizer, text_encoder, prompts, negative_prompt="", return_full_only=False, one_uncond_input_only=False):51    if negative_prompt == "":52        print("Note that negative_prompt is an empty string")53    54    text_input = tokenizer(55        prompts, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt"56    )57    58    max_length = text_input.input_ids.shape[-1]59    if one_uncond_input_only:60        num_uncond_input = 161    else:62        num_uncond_input = len(prompts)63    uncond_input = tokenizer([negative_prompt] * num_uncond_input, padding="max_length", max_length=max_length, return_tensors="pt")64 65    with torch.no_grad():66        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]67        cond_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]68    69    if one_uncond_input_only:70        return uncond_embeddings, cond_embeddings71    72    text_embeddings = torch.cat([uncond_embeddings, cond_embeddings])73    74    if return_full_only:75        return text_embeddings76    return text_embeddings, uncond_embeddings, cond_embeddings77 78def process_input_embeddings(input_embeddings):79    assert isinstance(input_embeddings, (tuple, list))80    if len(input_embeddings) == 3:81        # input_embeddings: text_embeddings, uncond_embeddings, cond_embeddings82        # Assume `uncond_embeddings` is full (has batch size the same as cond_embeddings)83        _, uncond_embeddings, cond_embeddings = input_embeddings84        assert uncond_embeddings.shape[0] == cond_embeddings.shape[0], f"{uncond_embeddings.shape[0]} != {cond_embeddings.shape[0]}"85        return input_embeddings86    elif len(input_embeddings) == 2:87        # input_embeddings: uncond_embeddings, cond_embeddings88        # uncond_embeddings may have only one item89        uncond_embeddings, cond_embeddings = input_embeddings90        if uncond_embeddings.shape[0] == 1:91            uncond_embeddings = uncond_embeddings.expand(cond_embeddings.shape)92        # We follow the convention: negative (unconditional) prompt comes first93        text_embeddings = torch.cat((uncond_embeddings, cond_embeddings), dim=0)94        return text_embeddings, uncond_embeddings, cond_embeddings95    else:96        raise ValueError(f"input_embeddings length: {len(input_embeddings)}")97