piyushgrover/Stable-Diffusion-Image-Generation

1from base64 import b64encode2 3import numpy4import torch5from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel6 7# For video display:8from matplotlib import pyplot as plt9from pathlib import Path10from PIL import Image11from torch import autocast12from torchvision import transforms as tfms13from tqdm.auto import tqdm14from transformers import CLIPTextModel, CLIPTokenizer, logging15import os16 17torch.manual_seed(1)18 19# Supress some unnecessary warnings when loading the CLIPTextModel20logging.set_verbosity_error()21 22# Set device23torch_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"24if "mps" == torch_device: os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = "1"25 26import gc27gc.collect()28torch.cuda.empty_cache()29 30from diffusers import StableDiffusionPipeline31 32model_id = "segmind/tiny-sd"33 34pipe = StableDiffusionPipeline.from_pretrained(model_id).to("cpu")35text_encoder = pipe.text_encoder.to(torch_device)36text_encoder.eval()37unet = pipe.unet.to(torch_device)38unet.eval()39vae = pipe.vae.to(torch_device)40vae.eval()41 42tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14')43scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)44del pipe45gc.collect()46 47seed_values = [0, 0, 0, 0, 0]48 49def load_learned_embeds():50    pathlist = Path('learned_embeds/').glob('*_learned_embeds.bin')51    learned_embeds = []52 53    for path in pathlist:54        path_in_str = str(path)55        # print(path_in_str)56        learned_embeds.append(torch.load(path_in_str))57 58    concept_embeds_list = []59    for obj in learned_embeds:60        for k, v in obj.items():61            if v.shape[0] == 768:62                print(k, v.shape)63                concept_embeds_list.append(v)64 65    return torch.stack(concept_embeds_list)66 67 68def pil_to_latent(input_im):69    # Single image -> single latent in a batch (so size 1, 4, 64, 64)70    with torch.no_grad():71        latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device) * 2 - 1)  # Note scaling72    return 0.18215 * latent.latent_dist.sample()73 74 75def latents_to_pil(latents):76    # bath of latents -> list of images77    latents = (1 / 0.18215) * latents78    with torch.no_grad():79        image = vae.decode(latents).sample80    image = (image / 2 + 0.5).clamp(0, 1)81    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()82    images = (image * 255).round().astype("uint8")83    pil_images = [Image.fromarray(image) for image in images]84    return pil_images85 86 87# Prep Scheduler88def set_timesteps(scheduler, num_inference_steps):89    scheduler.set_timesteps(num_inference_steps)90    scheduler.timesteps = scheduler.timesteps.to(91        torch.float32)  # minor fix to ensure MPS compatibility, fixed in diffusers PR 392592 93 94def get_output_embeds(input_embeddings):95    # CLIP's text model uses causal mask, so we prepare it here:96    bsz, seq_len = input_embeddings.shape[:2]97    causal_attention_mask = text_encoder.text_model._build_causal_attention_mask(bsz, seq_len,98                                                                                 dtype=input_embeddings.dtype)99 100    # Getting the output embeddings involves calling the model with passing output_hidden_states=True101    # so that it doesn't just return the pooled final predictions:102    encoder_outputs = text_encoder.text_model.encoder(103        inputs_embeds=input_embeddings,104        attention_mask=None,  # We aren't using an attention mask so that can be None105        causal_attention_mask=causal_attention_mask.to(torch_device),106        output_attentions=None,107        output_hidden_states=True,  # We want the output embs not the final output108        return_dict=None,109    )110 111    # We're interested in the output hidden state only112    output = encoder_outputs[0]113 114    # There is a final layer norm we need to pass these through115    output = text_encoder.text_model.final_layer_norm(output)116 117    # And now they're ready!118    return output119 120def blue_loss(images, contrast_perc=80):121    # How far the pixels are from +80% contrast:122    contrast = 255*contrast_perc // 100 # it ranges from -255 to +255123    contrast_scale_factor = (259 * (contrast + 255)) / (255 * (259 - contrast))124    cimgs = (contrast_scale_factor * (images - 0.5) + 0.5 )125    cimgs = torch.where(cimgs > 1.0, 1.0, cimgs)126    cimgs = torch.where(cimgs < 0.0, 0.0, cimgs)127    error = torch.abs( images - cimgs ).mean()128    #error = torch.abs(images[:] - 0.9).mean() # [:,2] -> all images in batch, only the blue channel129    print('error: ', error)130    return error131 132# Generating an image with these modified embeddings133def generate_with_embs(text_input, text_embeddings, output=None, generator=None, contrast_loss=False, contrast_perc=0):134    height = 512  # default height of Stable Diffusion135    width = 512  # default width of Stable Diffusion136    num_inference_steps = 30  # Number of denoising steps137    guidance_scale = 7.5  # Scale for classifier-free guidance138 139    if generator is None:140        generator = torch.manual_seed(32)  # Seed generator to create the inital latent noise141 142    batch_size = 1143 144    max_length = text_input.input_ids.shape[-1]145    uncond_input = tokenizer(146        [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"147    )148    with torch.no_grad():149        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]150    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])151 152    # Prep Scheduler153    set_timesteps(scheduler, num_inference_steps)154 155    # Prep latents156    latents = torch.randn(157        (batch_size, unet.in_channels, height // 8, width // 8),158        generator=generator,159    )160    latents = latents.to(torch_device)161    latents = latents * scheduler.init_noise_sigma162 163    # Loop164    #for i, t in tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps)):165    for i, t in enumerate(scheduler.timesteps):166        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.167        latent_model_input = torch.cat([latents] * 2)168        sigma = scheduler.sigmas[i]169        latent_model_input = scheduler.scale_model_input(latent_model_input, t)170 171        # predict the noise residual172        with torch.no_grad():173            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]174 175        # perform guidance176        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)177        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)178 179        #### ADDITIONAL GUIDANCE ###180        if contrast_loss:181            blue_loss_scale = 70182            if i % 5 == 0:183                # Requires grad on the latents184                latents = latents.detach().requires_grad_()185 186                # Get the predicted x0:187                latents_x0 = latents - sigma * noise_pred188                # latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample189 190                # Decode to image space191                denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5  # range (0, 1)192 193                # Calculate loss194                loss = blue_loss(denoised_images, contrast_perc=contrast_perc) * blue_loss_scale195 196                # Occasionally print it out197                if i % 10 == 0:198                    print(i, 'loss:', loss.item())199 200                # Get gradient201                cond_grad = torch.autograd.grad(loss, latents)[0]202 203                # Modify the latents based on this gradient204                latents = latents.detach() - cond_grad * sigma ** 2205 206        # compute the previous noisy sample x_t -> x_t-1207        latents = scheduler.step(noise_pred, t, latents).prev_sample208        if output:209            output = latents_to_pil(latents)[0]210 211    return latents_to_pil(latents)[0]212 213 214concept_embeds = load_learned_embeds()215 216token_emb_layer = text_encoder.text_model.embeddings.token_embedding217#token_emb_layer  # Vocab size 49408, emb_dim 768218 219pos_emb_layer = text_encoder.text_model.embeddings.position_embedding220#pos_emb_layer221 222def func_generate(query, concept_idx, seed_start, contrast_loss=False, contrast_perc=None):223    prompt = query + ' in the style of bulb'224    text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True,225                           return_tensors="pt")226    input_ids = text_input.input_ids.to(torch_device)227 228    # Get token embeddings229    position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]230    position_embeddings = pos_emb_layer(position_ids)231 232    s = seed_start233 234    token_embeddings = token_emb_layer(input_ids)235    # The new embedding - our special birb word236    replacement_token_embedding = concept_embeds[concept_idx].to(torch_device)237 238    # Insert this into the token embeddings239    token_embeddings[0, torch.where(input_ids[0] == 22373)] = replacement_token_embedding.to(torch_device)240 241    # Combine with pos embs242    input_embeddings = token_embeddings + position_embeddings243 244    #  Feed through to get final output embs245    modified_output_embeddings = get_output_embeds(input_embeddings)246 247    # And generate an image with this:248 249    if contrast_loss and seed_values[concept_idx] > 0:250        s = seed_values[concept_idx]251    else:252        s = random.randint(s + 1, s + 30)253        seed_values[concept_idx] = s254    255    g = torch.manual_seed(s)256    return generate_with_embs(text_input, modified_output_embeddings, generator=g, contrast_loss=contrast_loss, contrast_perc=contrast_perc)257