Upload 10 files

README.md

@@ -1,13 +1,89 @@
 ---
-title: Session20
-emoji: 🏢
-colorFrom: red
-colorTo: purple
+title: "ERA SESSION20 - Stable Diffusion: Generative Art with Guidance"
+emoji: 🌍
+colorFrom: indigo
+colorTo: pink
 sdk: gradio
-sdk_version: 3.47.1
+sdk_version: 3.48.0
 app_file: app.py
 pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+**Styles Used:**
+1. [Oil style](https://huggingface.co/sd-concepts-library/oil-style)
+2. [Xyz](https://huggingface.co/sd-concepts-library/xyz)
+3. [Allante](https://huggingface.co/sd-concepts-library/style-of-marc-allante)
+4. [Moebius](https://huggingface.co/sd-concepts-library/moebius)
+5. [Polygons](https://huggingface.co/sd-concepts-library/low-poly-hd-logos-icons)
+
+### Result of Experiments with different styles:
+**Prompt:** `"a cat and dog in the style of cs"` \
+_"cs" in the prompt refers to "custom style" whose embedding is replaced by each of the concept embeddings shown below_
+![image](https://github.com/RaviNaik/ERA-SESSION20/assets/23289802/1effe375-6ef4-4adc-be7b-d6311fdaa50d)
+
+---
+**Prompt:** `"dolphin swimming on Mars in the style of cs"`
+![image](https://github.com/RaviNaik/ERA-SESSION20/assets/23289802/2cd32248-4233-42c0-97c0-00e1ae8fdc85)
+
+### Result of Experiments with Guidance loss functions:
+**Prompt:** `"a mouse in the style of cs"`
+**Loss Function:**
+```python
+def loss_fn(images):
+    return images.mean()
+```
+![image](https://github.com/RaviNaik/ERA-SESSION20/assets/23289802/c9d46e14-44bb-4ea7-88a4-26ef46344fce)
+---
+```python
+def loss_fn(images):
+    return -images.median()/3
+```
+![image](https://github.com/RaviNaik/ERA-SESSION20/assets/23289802/2649e4f6-3de5-4e54-8f22-3d65874b7b07)
+---
+```python
+def loss_fn(images):
+    error = (images - images.min()) / 255*(images.max() - images.min())
+    return error.mean()
+```
+![image](https://github.com/RaviNaik/ERA-SESSION20/assets/23289802/6399c780-e9b7-42f8-8d90-44c8b40d5265)
+---
+**Prompt:** `"angry german shephard in the style of cs"`
+```python
+def loss_fn(images):
+    error1 = torch.abs(images[:, 0] - 0.9)
+    error2 = torch.abs(images[:, 1] - 0.9)
+    error3 = torch.abs(images[:, 2] - 0.9)
+    return (
+        torch.sin(error1.mean()) + torch.sin(error2.mean()) + torch.sin(error3.mean())
+    ) / 3
+```
+![image](https://github.com/RaviNaik/ERA-SESSION20/assets/23289802/fa7d30ed-4efd-4504-b89c-94e093f51f9c)
+
+---
+**Prompt:** `"A campfire (oil on canvas)"`
+```python
+def loss_fn(images):
+    error1 = torch.abs(images[:, 0] - 0.9)
+    error2 = torch.abs(images[:, 1] - 0.9)
+    error3 = torch.abs(images[:, 2] - 0.9)
+    return (
+        torch.sin((error1 * error2 * error3)).mean()
+        + torch.cos((error1 * error2 * error3)).mean()
+    )
+```
+![image](https://github.com/RaviNaik/ERA-SESSION20/assets/23289802/88382dae-6701-4103-a664-ed17727b690f)
+
+---
+```python
+def loss_fn(images):
+    error1 = torch.abs(images[:, 0] - 0.9)
+    error2 = torch.abs(images[:, 1] - 0.9)
+    error3 = torch.abs(images[:, 2] - 0.9)
+    return (
+        torch.sin(error1.mean()) + torch.sin(error2.mean()) + torch.sin(error3.mean())
+    ) / 3
+```
+![image](https://github.com/RaviNaik/ERA-SESSION20/assets/23289802/0ab3edad-579d-4821-b992-6c18b61bd444)
+
+

app.py

@@ -1,307 +1,91 @@
 import gradio as gr
-from base64 import b64encode
-import numpy
+import random
 import torch
-from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
-from PIL import Image
-from torch import autocast
-from torchvision import transforms as tfms
-from tqdm.auto import tqdm
-from transformers import CLIPTextModel, CLIPTokenizer, logging
-import torchvision.transforms as T
-
-torch.manual_seed(1)
-logging.set_verbosity_error()
-torch_device = "cpu"
-
-vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
-
-# Load the tokenizer and text encoder to tokenize and encode the text.
-tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
-text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
-
-# The UNet model for generating the latents.
-unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
-
-# The noise scheduler
-scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
-
-vae = vae.to(torch_device)
-text_encoder = text_encoder.to(torch_device)
-unet = unet.to(torch_device);
-
-token_emb_layer = text_encoder.text_model.embeddings.token_embedding
-pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
-position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]
-position_embeddings = pos_emb_layer(position_ids)
-
-def pil_to_latent(input_im):
-    # Single image -> single latent in a batch (so size 1, 4, 64, 64)
-    with torch.no_grad():
-        latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling
-    return 0.18215 * latent.latent_dist.sample()
-
-def latents_to_pil(latents):
-    # bath of latents -> list of images
-    latents = (1 / 0.18215) * latents
-    with torch.no_grad():
-        image = vae.decode(latents).sample
-    image = (image / 2 + 0.5).clamp(0, 1)
-    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
-    images = (image * 255).round().astype("uint8")
-    pil_images = [Image.fromarray(image) for image in images]
-    return pil_images
-
-def get_output_embeds(input_embeddings):
-    # CLIP's text model uses causal mask, so we prepare it here:
-    bsz, seq_len = input_embeddings.shape[:2]
-    causal_attention_mask = text_encoder.text_model._build_causal_attention_mask(bsz, seq_len, dtype=input_embeddings.dtype)
-
-    # Getting the output embeddings involves calling the model with passing output_hidden_states=True
-    # so that it doesn't just return the pooled final predictions:
-    encoder_outputs = text_encoder.text_model.encoder(
-        inputs_embeds=input_embeddings,
-        attention_mask=None, # We aren't using an attention mask so that can be None
-        causal_attention_mask=causal_attention_mask.to(torch_device),
-        output_attentions=None,
-        output_hidden_states=True, # We want the output embs not the final output
-        return_dict=None,
-    )
-
-    # We're interested in the output hidden state only
-    output = encoder_outputs[0]
-
-    # There is a final layer norm we need to pass these through
-    output = text_encoder.text_model.final_layer_norm(output)
-
-    # And now they're ready!
-    return output
-
-def generate_with_embs(text_embeddings, seed, max_length):
-    height = 512                        # default height of Stable Diffusion
-    width = 512                         # default width of Stable Diffusion
-    num_inference_steps = 10            # Number of denoising steps
-    guidance_scale = 7.5                # Scale for classifier-free guidance
-    generator = torch.manual_seed(seed)   # Seed generator to create the inital latent noise
-    batch_size = 1
-
-    # max_length = text_input.input_ids.shape[-1]
-    uncond_input = tokenizer(
-      [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
-    )
-    with torch.no_grad():
-        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
-    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
-
-    # Prep Scheduler
-    set_timesteps(scheduler, num_inference_steps)
-
-    # Prep latents
-    latents = torch.randn(
-    (batch_size, unet.in_channels, height // 8, width // 8),
-    generator=generator,
-    )
-    latents = latents.to(torch_device)
-    latents = latents * scheduler.init_noise_sigma
-
-    # Loop
-    for i, t in tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps)):
-        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
-        latent_model_input = torch.cat([latents] * 2)
-        sigma = scheduler.sigmas[i]
-        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
-
-        # predict the noise residual
-        with torch.no_grad():
-            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]
-
-        # perform guidance
-        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-        # compute the previous noisy sample x_t -> x_t-1
-        latents = scheduler.step(noise_pred, t, latents).prev_sample
-
-    return latents_to_pil(latents)[0]
-
-# Prep Scheduler
-def set_timesteps(scheduler, num_inference_steps):
-    scheduler.set_timesteps(num_inference_steps)
-    scheduler.timesteps = scheduler.timesteps.to(torch.float32) # minor fix to ensure MPS compatibility, fixed in diffusers PR 3925
-
-def embed_style(prompt, style_embed, style_seed):
-    # Tokenize
-    text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
-    input_ids = text_input.input_ids.to(torch_device)
-
-    # Get token embeddings
-    token_embeddings = token_emb_layer(input_ids)
-
-    replacement_token_embedding = style_embed.to(torch_device)
-
-    # Insert this into the token embeddings
-    token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)
-
-    # Combine with pos embs
-    input_embeddings = token_embeddings + position_embeddings
-
-    #  Feed through to get final output embs
-    modified_output_embeddings = get_output_embeds(input_embeddings)
-
-    # And generate an image with this:
-    max_length = text_input.input_ids.shape[-1]
-    return generate_with_embs(modified_output_embeddings, style_seed, max_length)
-
-def loss_style(prompt, style_embed, style_seed):
-    # Tokenize
-    text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
-    input_ids = text_input.input_ids.to(torch_device)
-
-    # Get token embeddings
-    token_embeddings = token_emb_layer(input_ids)
-
-    # The new embedding - our special birb word
-    replacement_token_embedding = style_embed.to(torch_device)
-
-    # Insert this into the token embeddings
-    token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)
-
-    # Combine with pos embs
-    input_embeddings = token_embeddings + position_embeddings
-
-    #  Feed through to get final output embs
-    modified_output_embeddings = get_output_embeds(input_embeddings)
-
-    # And generate an image with this:
-    max_length = text_input.input_ids.shape[-1]
-    return generate_loss_based_image(modified_output_embeddings, style_seed,max_length)
-
-
-def color_loss(image):
-    color_channel = image[:, 1] 
-    target_value = 0.7 
-    error = torch.abs(color_channel - target_value).mean()
-    return error
-
-def generate_loss_based_image(text_embeddings, seed, max_length):
-
-    height = 64               
-    width = 64                    
-    num_inference_steps = 10       
-    guidance_scale = 8             
-    generator = torch.manual_seed(64)  
-    batch_size = 1
-    loss_scale = 200
-
-    uncond_input = tokenizer(
-        [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
-    )
-    with torch.no_grad():
-        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
-    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
-
-    # Prep Scheduler
-    set_timesteps(scheduler, num_inference_steps+1)
-
-    # Prep latents
-    latents = torch.randn(
-    (batch_size, unet.in_channels, height // 8, width // 8),
-    generator=generator,
-    )
-    latents = latents.to(torch_device)
-    latents = latents * scheduler.init_noise_sigma
-
-    sched_out = None
-
-    # Loop
-    for i, t in tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps)):
-        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
-        latent_model_input = torch.cat([latents] * 2)
-        sigma = scheduler.sigmas[i]
-        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
-
-        # predict the noise residual
-        with torch.no_grad():
-            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]
-
-        # perform CFG
-        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-        ### ADDITIONAL GUIDANCE ###
-        if i%5 == 0 and i>0:
-            # Requires grad on the latents
-            latents = latents.detach().requires_grad_()
-
-            # Get the predicted x0:
-            scheduler._step_index -= 1
-            latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample
-
-            # Decode to image space
-            denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5 # range (0, 1)
-
-
-            # Calculate loss
-            loss = color_loss(denoised_images) * loss_scale
-
-            # Occasionally print it out
-            # if i%10==0:
-            print(i, 'loss:', loss)
-
-            # Get gradient
-            cond_grad = torch.autograd.grad(loss, latents)[0]
-
-            # Modify the latents based on this gradient
-            latents = latents.detach() - cond_grad * sigma**2
-            # To PIL Images
-            im_t0 = latents_to_pil(latents_x0)[0]
-            im_next = latents_to_pil(latents)[0]
-
-        # Now step with scheduler
-        latents = scheduler.step(noise_pred, t, latents).prev_sample
-        
-    return latents_to_pil(latents)[0]
-
-
-def generate_image_from_prompt(text_in, style_in):
-    STYLE_LIST = ['coffeemachine.bin', 'collage_style.bin', 'cube.bin', 'jerrymouse2.bin', 'zero.bin']
-    STYLE_SEEDS = [32, 64, 128, 16, 8]
-    
-    print(text_in)
-    print(style_in)
-    style_file = style_in + '.bin'
-    idx = STYLE_LIST.index(style_file)
-    print(style_file)
-    print(idx)  
-    
-    prompt = text_in + ' a puppy'
-    
-    style_seed = STYLE_SEEDS[idx]
-    style_dict = torch.load(style_file)
-    style_embed = [v for v in style_dict.values()]
-
-    generated_image = embed_style(prompt, style_embed[0], style_seed)
-
-    loss_generated_img = (loss_style(prompt, style_embed[0], style_seed))
-
-    return [generated_image, loss_generated_img]
-    
-
-# Define Interface
-
-title = 'ERA-SESSION20 Generative Art and Stable Diffusion'
-
-demo = gr.Interface(generate_image_from_prompt,
-                    inputs = [gr.Textbox(1, label='prompt'),
-                              gr.Dropdown(
-                                ['coffeemachine', 'collage_style', 'cube', 'jerrymouse2', 'zero'],value="cube", label="Pretrained Styles"
-                            )
-                             ],
-                    outputs = [
-                                
-                        gr.Gallery(label="Generated images", show_label=True, elem_id="gallery", columns=[2], rows=[2], object_fit="contain", height="auto")
-                              ],
-                    
-                    title = title
-                   )
-demo.launch(debug=True)
+import pathlib
+
+from src.utils import concept_styles, loss_fn
+from src.stable_diffusion import StableDiffusion
+
+PROJECT_PATH = "."
+CONCEPT_LIBS_PATH = f"{PROJECT_PATH}/concept_libs"
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def generate(prompt, styles, gen_steps, loss_scale):
+    lossless_images, lossy_images = [], []
+    for style in styles:
+        concept_lib_path = f"{CONCEPT_LIBS_PATH}/{concept_styles[style]}"
+        concept_lib = pathlib.Path(concept_lib_path)
+        concept_embed = torch.load(concept_lib)
+
+        manual_seed = random.randint(0, 100)
+        diffusion = StableDiffusion(
+            device=DEVICE,
+            num_inference_steps=gen_steps,
+            manual_seed=manual_seed,
+        )
+        generated_image_lossless = diffusion.generate_image(
+            prompt=prompt,
+            loss_fn=loss_fn,
+            loss_scale=0,
+            concept_embed=concept_embed,
+        )
+        generated_image_lossy = diffusion.generate_image(
+            prompt=prompt,
+            loss_fn=loss_fn,
+            loss_scale=loss_scale,
+            concept_embed=concept_embed,
+        )
+        lossless_images.append((generated_image_lossless, style))
+        lossy_images.append((generated_image_lossy, style))
+    return {lossless_gallery: lossless_images, lossy_gallery: lossy_images}
+
+
+with gr.Blocks() as app:
+    gr.Markdown("## ERA Session20 - Stable Diffusion: Generative Art with Guidance")
+    with gr.Row():
+        with gr.Column():
+            prompt_box = gr.Textbox(label="Prompt", interactive=True)
+            style_selector = gr.Dropdown(
+                choices=list(concept_styles.keys()),
+                value=list(concept_styles.keys())[0],
+                multiselect=True,
+                label="Select a Concept Style",
+                interactive=True,
+            )
+            gen_steps = gr.Slider(
+                minimum=10,
+                maximum=50,
+                value=30,
+                step=10,
+                label="Select Number of Steps",
+                interactive=True,
+            )
+
+            loss_scale = gr.Slider(
+                minimum=0,
+                maximum=32,
+                value=8,
+                step=8,
+                label="Select Guidance Scale",
+                interactive=True,
+            )
+
+            submit_btn = gr.Button(value="Generate")
+
+        with gr.Column():
+            lossless_gallery = gr.Gallery(
+                label="Generated Images without Guidance", show_label=True
+            )
+            lossy_gallery = gr.Gallery(
+                label="Generated Images with Guidance", show_label=True
+            )
+
+        submit_btn.click(
+            generate,
+            inputs=[prompt_box, style_selector, gen_steps, loss_scale],
+            outputs=[lossless_gallery, lossy_gallery],
+        )
+
+app.launch()

concept_libs/coffeemachine.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc3a85dc9cbdf6ab5fca4056c473da1b632c0565030be918682ce3e62095b4b1
+size 3840

concept_libs/collage_style.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b143c4841c5f2d39d0eb2015d62c17d1b18da9bb0a42c76320df7acfe1e144bf
+size 3840

concept_libs/cube.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a6d6394f0cd38847259c42746a6b0e50ca1e76e6ddc8e217ff14f2feb7dbca4
+size 3819

concept_libs/jerrymouse2.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9713d9367f1faa6ebd753db5c8a209c565be0b25e32051c723c4533dd9df605
+size 3840

concept_libs/zero.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:78286aa910deafe4e46c6e38a86f464a246aef95ad5611a756dd99405f418a85
+size 3819

src/stable_diffusion.py

@@ -0,0 +1,222 @@
+import torch
+from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
+from transformers import CLIPTextModel, CLIPTokenizer
+from PIL import Image
+from tqdm import tqdm
+
+
+class StableDiffusion:
+    def __init__(
+        self,
+        vae_arch="CompVis/stable-diffusion-v1-4",
+        tokenizer_arch="openai/clip-vit-large-patch14",
+        encoder_arch="openai/clip-vit-large-patch14",
+        unet_arch="CompVis/stable-diffusion-v1-4",
+        device="cpu",
+        height=512,
+        width=512,
+        num_inference_steps=30,
+        guidance_scale=7.5,
+        manual_seed=1,
+    ) -> None:
+        self.height = height  # default height of Stable Diffusion
+        self.width = width  # default width of Stable Diffusion
+        self.num_inference_steps = num_inference_steps  # Number of denoising steps
+        self.guidance_scale = guidance_scale  # Scale for classifier-free guidance
+        self.device = device
+        self.manual_seed = manual_seed
+
+        vae = AutoencoderKL.from_pretrained(vae_arch, subfolder="vae")
+        # Load the tokenizer and text encoder to tokenize and encode the text.
+        self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer_arch)
+        text_encoder = CLIPTextModel.from_pretrained(encoder_arch)
+
+        # The UNet model for generating the latents.
+        unet = UNet2DConditionModel.from_pretrained(unet_arch, subfolder="unet")
+
+        # The noise scheduler
+        self.scheduler = LMSDiscreteScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            num_train_timesteps=1000,
+        )
+
+        # To the GPU we go!
+        self.vae = vae.to(self.device)
+        self.text_encoder = text_encoder.to(self.device)
+        self.unet = unet.to(self.device)
+
+        self.token_emb_layer = text_encoder.text_model.embeddings.token_embedding
+        pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
+        position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]
+        self.position_embeddings = pos_emb_layer(position_ids)
+
+    def get_output_embeds(self, input_embeddings):
+        # CLIP's text model uses causal mask, so we prepare it here:
+        bsz, seq_len = input_embeddings.shape[:2]
+        causal_attention_mask = (
+            self.text_encoder.text_model._build_causal_attention_mask(
+                bsz, seq_len, dtype=input_embeddings.dtype
+            )
+        )
+
+        # Getting the output embeddings involves calling the model with passing output_hidden_states=True
+        # so that it doesn't just return the pooled final predictions:
+        encoder_outputs = self.text_encoder.text_model.encoder(
+            inputs_embeds=input_embeddings,
+            attention_mask=None,  # We aren't using an attention mask so that can be None
+            causal_attention_mask=causal_attention_mask.to(self.device),
+            output_attentions=None,
+            output_hidden_states=True,  # We want the output embs not the final output
+            return_dict=None,
+        )
+
+        # We're interested in the output hidden state only
+        output = encoder_outputs[0]
+
+        # There is a final layer norm we need to pass these through
+        output = self.text_encoder.text_model.final_layer_norm(output)
+
+        # And now they're ready!
+        return output
+
+    def set_timesteps(self, scheduler, num_inference_steps):
+        scheduler.set_timesteps(num_inference_steps)
+        scheduler.timesteps = scheduler.timesteps.to(torch.float32)
+
+    def latents_to_pil(self, latents):
+        # bath of latents -> list of images
+        latents = (1 / 0.18215) * latents
+        with torch.no_grad():
+            image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
+        images = (image * 255).round().astype("uint8")
+        pil_images = [Image.fromarray(image) for image in images]
+        return pil_images
+
+    def generate_with_embs(self, text_embeddings, text_input, loss_fn, loss_scale):
+        generator = torch.manual_seed(
+            self.manual_seed
+        )  # Seed generator to create the inital latent noise
+        batch_size = 1
+
+        max_length = text_input.input_ids.shape[-1]
+        uncond_input = self.tokenizer(
+            [""] * batch_size,
+            padding="max_length",
+            max_length=max_length,
+            return_tensors="pt",
+        )
+        with torch.no_grad():
+            uncond_embeddings = self.text_encoder(
+                uncond_input.input_ids.to(self.device)
+            )[0]
+        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+
+        # Prep Scheduler
+        self.set_timesteps(self.scheduler, self.num_inference_steps)
+
+        # Prep latents
+        latents = torch.randn(
+            (batch_size, self.unet.in_channels, self.height // 8, self.width // 8),
+            generator=generator,
+        )
+        latents = latents.to(self.device)
+        latents = latents * self.scheduler.init_noise_sigma
+
+        # Loop
+        for i, t in tqdm(
+            enumerate(self.scheduler.timesteps), total=len(self.scheduler.timesteps)
+        ):
+            # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+            latent_model_input = torch.cat([latents] * 2)
+            sigma = self.scheduler.sigmas[i]
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            # predict the noise residual
+            with torch.no_grad():
+                noise_pred = self.unet(
+                    latent_model_input, t, encoder_hidden_states=text_embeddings
+                )["sample"]
+
+            # perform guidance
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + self.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+            if i % 5 == 0:
+                # Requires grad on the latents
+                latents = latents.detach().requires_grad_()
+
+                # Get the predicted x0:
+                # latents_x0 = latents - sigma * noise_pred
+                latents_x0 = self.scheduler.step(
+                    noise_pred, t, latents
+                ).pred_original_sample
+
+                # Decode to image space
+                denoised_images = (
+                    self.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
+                )  # range (0, 1)
+
+                # Calculate loss
+                loss = loss_fn(denoised_images) * loss_scale
+
+                # Occasionally print it out
+                # if i % 10 == 0:
+                #     print(i, "loss:", loss.item())
+
+                # Get gradient
+                cond_grad = torch.autograd.grad(loss, latents)[0]
+
+                # Modify the latents based on this gradient
+                latents = latents.detach() - cond_grad * sigma**2
+                self.scheduler._step_index = self.scheduler._step_index - 1
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        return self.latents_to_pil(latents)[0]
+
+    def generate_image(
+        self,
+        prompt="A campfire (oil on canvas)",
+        loss_fn=None,
+        loss_scale=200,
+        concept_embed=None,  # birb_embed["<birb-style>"]
+    ):
+        prompt += " in the style of cs"
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        input_ids = text_input.input_ids.to(self.device)
+        custom_style_token = self.tokenizer.encode("cs", add_special_tokens=False)[0]
+        # Get token embeddings
+        token_embeddings = self.token_emb_layer(input_ids)
+
+        # The new embedding - our special birb word
+        embed_key = list(concept_embed.keys())[0]
+        replacement_token_embedding = concept_embed[embed_key]
+
+        # Insert this into the token embeddings
+        token_embeddings[
+            0, torch.where(input_ids[0] == custom_style_token)
+        ] = replacement_token_embedding.to(self.device)
+        # token_embeddings = token_embeddings + (replacement_token_embedding * 0.9)
+        # Combine with pos embs
+        input_embeddings = token_embeddings + self.position_embeddings
+
+        #  Feed through to get final output embs
+        modified_output_embeddings = self.get_output_embeds(input_embeddings)
+
+        # And generate an image with this:
+        generated_image = self.generate_with_embs(
+            modified_output_embeddings, text_input, loss_fn, loss_scale
+        )
+        return generated_image

src/utils.py

@@ -0,0 +1,11 @@
+def loss_fn(images):
+    return -images.median() / 3
+
+
+concept_styles = {
+    "Coffee Machine": "coffeemachine.bin",
+    "College Style": "college_style.bin",
+    "Cube": "cube.bin",
+    "Jerry Mouse": "jerrymouse",
+    "Zero": "zero.bin",
+}