Upload folder using huggingface_hub

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+from sd import stl_list, img_size_opt_dict, generate_images
+
+with gr.Blocks() as demo:
+    gr.HTML("<h1 align = 'center'> Stable Diffusion - Text Inversion and additional guidence</h1>")
+    gr.HTML("<h4 align = 'center'> Generates imgaes based on the prompt and 5 different styles and then with additional guidence of hue loss</h4>")
+    gr.HTML("<h6 align = 'center'> !!The image generation may take 5 to 10 minutes on CPU!!</h4>")
+
+    with gr.Row():
+        content = gr.Textbox(label = "Enter prompt text here")
+        gr.Examples([
+            "A mouse",
+            "A puppy"
+        ],
+        inputs = content)
+        num_steps = gr.Slider(1, 50, step = 1, value=30, label="Number of inference steps", info="Choose between 1 and 50")
+        # gr.Number(value = 10, label = "Number of inference steps")
+
+    
+    with gr.Row():
+        stl_dropdown = gr.Dropdown(
+            stl_list, 
+            value=stl_list[:1], multiselect=True, label="Style", 
+            info="Styles to be applied on images"
+        )
+        size_dropdown = gr.Dropdown(
+            [*img_size_opt_dict], 
+            value = [*img_size_opt_dict][-1],
+            label="Image size", info="Target size for generated images"
+        )
+    
+    inputs = [
+                content,
+                num_steps,
+                stl_dropdown,
+                size_dropdown
+                ]
+    
+    generate_btn = gr.Button(value = 'Generate')
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            wo_add_guide = gr.Gallery(
+            label="Without additional guidence", show_label=True, elem_id="gallery"
+            , columns=[3], rows=[2], object_fit="contain", height="auto")
+
+        with gr.Column(scale=2):
+            add_guide = gr.Gallery(
+            label="With hue loss guidence", show_label=True, elem_id="gallery"
+            , columns=[3], rows=[2], object_fit="contain", height="auto")
+        outputs  = [wo_add_guide, add_guide ]
+    generate_btn.click(fn = generate_images, inputs= inputs, outputs = outputs)
+
+# for collab
+# demo.launch(debug=True)
+
+if __name__ == '__main__':
+    demo.launch()

hue_loss.py

@@ -0,0 +1,41 @@
+import torch
+
+# hue loss
+def rgb_to_hsv(image):
+    r, g, b = image[:, 0, :, :], image[:, 1, :, :], image[:, 2, :, :]
+    maxc = torch.max(image, dim=1)[0]
+    minc = torch.min(image, dim=1)[0]
+
+    v = maxc
+    s = (maxc - minc) / (maxc + 1e-10)
+    deltac = maxc - minc
+
+    # Initialize hue
+    h = torch.zeros_like(maxc)
+
+    mask = maxc == r
+    h[mask] = ((g - b) / deltac)[mask] % 6
+
+    mask = maxc == g
+    h[mask] = ((b - r) / deltac)[mask] + 2
+
+    mask = maxc == b
+    h[mask] = ((r - g) / deltac)[mask] + 4
+
+    h = h / 6  # Normalize to [0, 1]
+    h[deltac == 0] = 0  # If no color difference, set hue to 0
+
+    return torch.stack([h, s, v], dim=1)
+
+
+def hue_loss(images, target_hue=0.5):
+    # Convert the images to HSV color space
+    hsv_images = rgb_to_hsv(images)
+
+    # Extract the hue channel
+    hue = hsv_images[:, 0, :, :]
+
+    # Calculate the error as the mean absolute deviation from the target hue
+    error = torch.abs(hue - target_hue).mean()
+
+    return error

requirements.txt

@@ -0,0 +1,6 @@
+pillow
+torch
+transformers
+diffusers
+huggingface_hub
+numpy

sd.py

@@ -0,0 +1,238 @@
+import numpy
+import torch
+from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
+
+from PIL import Image
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+import os
+
+from hue_loss import hue_loss
+
+
+torch.manual_seed(1)
+# if not (Path.home()/'.cache/huggingface'/'token').exists(): notebook_login()
+
+# Supress some unnecessary warnings when loading the CLIPTextModel
+logging.set_verbosity_error()
+
+# Set device
+torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+from huggingface_hub import hf_hub_download
+
+
+stl_list = [
+            'birb-style',
+            'cute-game-style',
+            'depthmap',
+            'line-art',
+            'low-poly-hd-logos-icons'
+            ]
+
+for stl in stl_list:
+    if not os.path.exists(stl):
+        os.mkdir(stl)
+    hf_hub_download(repo_id=f"sd-concepts-library/{stl}", filename="learned_embeds.bin", local_dir=f"./{stl}")
+
+img_size_opt_dict = {
+    "512x512 - best quality but very slow": (512,512), 
+    "256x256 - not good quality but still slow" :  (256,256), 
+    "128x128 - poor quality  but faster" : (128,128),
+    }
+
+# Load the autoencoder model which will be used to decode the latents into image space.
+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)
+
+# To the GPU we go!
+vae = vae.to(torch_device)
+text_encoder = text_encoder.to(torch_device)
+unet = unet.to(torch_device);
+
+# Convert latents to images
+
+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
+
+# 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
+
+#Generating an image with these modified embeddings
+def generate_with_embs(text_embeddings, text_input, loss_fn = None, loss_scale = 200, guidance_scale = 7.5,
+                       seed_value = 1, num_inference_steps = 50, additional_guidence = False, hight_width = (512, 512)):
+    height, width = hight_width                        # default height of Stable Diffusion
+    # width = 512                         # default width of Stable Diffusion
+    # num_inference_steps = 50            # Number of denoising steps
+                   # Scale for classifier-free guidance
+    generator = torch.manual_seed(seed_value)   # 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)
+
+        #### ADDITIONAL GUIDANCE ###
+        if i%5 == 0 and additional_guidence:
+            # Requires grad on the latents
+            latents = latents.detach().requires_grad_()
+
+            # Get the predicted x0:
+            latents_x0 = latents - sigma * noise_pred
+            # 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 = 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
+            latents = latents.detach() - cond_grad * sigma**2
+
+        # compute the previous noisy sample x_t -> x_t-1
+        latents = scheduler.step(noise_pred, t, latents).prev_sample
+
+        # Ensure the latents do not lose the grad tracking
+        # latents.requires_grad_()
+
+    return latents_to_pil(latents)[0]
+
+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
+
+# Access the embedding layer
+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 generate_images(prompt, num_inference_steps, stl_list, img_size):
+    ### add a statis text that will contain the style
+    prompt = prompt + ' in the style of puppy'
+    height_width = img_size_opt_dict[img_size]
+    # 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)
+
+    wo_guide_lst = []
+    guide_lst = []
+    for i, stl in enumerate(stl_list):
+        stl_embed = torch.load(f'{stl}/learned_embeds.bin')
+
+        # The new embedding - our special birb word
+        replacement_token_embedding = stl_embed[f'<{stl}>'].to(torch_device)
+
+        # Insert this into the token embeddings
+        token_embeddings[0, min(torch.where(input_ids[0]==tokenizer.eos_token_id)[0]) - 1] = 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:
+        pil_im = generate_with_embs(modified_output_embeddings,
+                                    num_inference_steps = num_inference_steps,
+                                    text_input = text_input,
+                                    seed_value = i,additional_guidence = False,
+                                    hight_width = height_width)
+        wo_guide_lst.append((pil_im,stl))
+
+        pil_im = generate_with_embs(modified_output_embeddings,
+                                    num_inference_steps = num_inference_steps,
+                                    text_input = text_input,
+                                    loss_fn = hue_loss,
+                                    additional_guidence = True,
+                                    hight_width = height_width,
+                                    seed_value = i)
+        guide_lst.append((pil_im,stl))
+
+    return wo_guide_lst, guide_lst