Initial Commit

app.py

@@ -0,0 +1,222 @@
+import gradio as gr
+from torchvision import transforms
+import torch
+from main_inference import generate_mixed_image, generate_image, progress_video
+import matplotlib.colors as mcolors
+
+def run_generate_mixed_image(prompt1,prompt2):
+    image = generate_mixed_image(prompt1,prompt2)
+    return image
+
+def run_generate_image(prompt1,noise_checkbox):
+    image = generate_image(prompt1,noised_image=noise_checkbox)
+    return image
+
+def run_generate_image_with_color_doninance(prompt1,color,color_loss_scale,noised_image_checkbox_1):
+    # Convert the hexadecimal color code to RGB values
+    rgba_color = mcolors.hex2color(color)
+    # Multiply the RGB values by 255 to get them in the [0, 255] range
+    rgb_values = [int(val * 255) for val in rgba_color]
+    image = generate_image(prompt1,True,color,color_loss_scale,noised_image_checkbox_1)
+    return image
+    
+def run_process_video(prompt):
+    # Ask for text input
+    video = progress_video(prompt)
+    return video
+
+
+description_text_to_image = """ ### Text to Image Generation
+                
+                1. Write a Text Prompt.
+                
+                2. Output will be an image based on the text prompt provided.
+                
+                3. Check if you want to see noised version of the image
+                                     
+              """
+
+description_generate_mixed_image = """ ### Mix Image Generation
+                
+                1. Write Two Text prompts.
+                
+                2. Output will a image which is mix of both of the text provided.
+                
+                3. Check if you want to see noised version of the image
+                                   
+              """
+
+description_generate_image_with_color_dominance = """ ### Generate Images with color dominance
+                
+                1. Write a Text Prompt.
+                
+                2. Select a color 
+                
+                3. Choose Color loss value
+                
+                4. Get the generated Image
+                
+                5. Check if you want to see noised version of the image
+                                     
+              """
+              
+description_progress_video = """ ### Get the full generation process video
+                
+                1. Write a Text Prompt.
+                
+                2. Output will be the video which contains frames of generated image, during various inference steps
+                                     
+              """
+
+# Description
+title = "<center><strong><font size='8'>The Stable Diffusion</font></strong></center>"
+
+image_input1 = gr.Image(type='pil') 
+image_input2 = gr.Image(type='filepath') 
+image_input3 = gr.Image(type='pil')
+image_input4 = gr.Image(type='pil')
+text_input = gr.Text(label="Enter Text Prompt")
+text_input2 = gr.Text(label="Enter Text Prompt")
+text_input3 = gr.Text(label="Enter Text Prompt")
+text_input4 = gr.Text(label = "Enter Text Prompt")
+text_input5 = gr.Text(label = "Enter Text Prompt")
+video_output = gr.Video()
+
+color = gr.ColorPicker(label="Select a Color",description="Choose a color from the color picker:")
+noised_image_checkbox  = gr.inputs.Checkbox(default=False, label="Show Noised Image")
+noised_image_checkbox_1  = gr.inputs.Checkbox(default=False, label="Show Noised Image")
+noised_image_checkbox_2  = gr.inputs.Checkbox(default=False, label="Show Noised Image")
+color_loss_scale = gr.inputs.Slider(minimum=0, maximum=255, default=40, step=1,label="Color Loss")
+css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
+
+with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Title
+            gr.Markdown(title)
+            
+    with gr.Tab("Generate Image"):
+        # Images
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                text_input.render()
+                noised_image_checkbox.render()
+            with gr.Column(scale=1):
+                image_input1.render()
+
+        # Submit & Clear
+        with gr.Row():
+            with gr.Column():
+                run_generate_image_button = gr.Button("generate_image", variant='primary')
+                clear_btn_text_to_image = gr.Button("Clear", variant="secondary")
+                gr.Markdown(description_text_to_image)
+                # gr.Examples(examples = ["A White cat", "a dog playing in garden", "people enjoying around sea"],
+                #             inputs=[text_input,red_color_slider,green_color_slider,blue_color_slider,noised_image_checkbox],
+                #             outputs=image_input1,
+                #             fn=run_generate_image,
+                #             cache_examples=True,
+                #             examples_per_page=4)
+
+    run_generate_image_button.click(run_generate_image,
+                        inputs=[text_input,noised_image_checkbox],
+                        outputs=image_input1)
+
+    with gr.Tab("Generate Image with Color Dominance"):
+        # Images
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                text_input4.render()
+                color_loss_scale.render()
+                noised_image_checkbox_1.render()
+                color.render()
+            with gr.Column(scale=1):
+                image_input3.render()
+
+        # Submit & Clear
+        with gr.Row():
+            with gr.Column():
+                run_generate_image_with_color_doninance_button = gr.Button("generate_image_with_color_doninance", variant='primary')
+                clear_btn_text_to_image = gr.Button("Clear", variant="secondary")
+                gr.Markdown(description_generate_image_with_color_dominance)
+                # gr.Examples(examples = ["A White cat", "a dog playing in garden", "people enjoying around sea"],
+                #             inputs=[text_input4,red_color_slider,green_color_slider,blue_color_slider,noised_image_checkbox],
+                #             outputs=image_input3,
+                #             fn=run_generate_image_with_color_doninance,
+                #             cache_examples=True,
+                #             examples_per_page=4)
+
+    run_generate_image_with_color_doninance_button.click(run_generate_image_with_color_doninance,
+                        inputs=[text_input4,color,color_loss_scale,noised_image_checkbox_1],
+                        outputs=image_input3)
+
+    ####################################################################################################################
+    with gr.Tab("Generate Mixed Image"):
+        # Images
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                text_input2.render()
+                text_input3.render()
+                noised_image_checkbox_2.render()
+            with gr.Column(scale=1):
+                image_input4.render()
+
+        # Submit & Clear
+        with gr.Row():
+            with gr.Column():
+                run_generate_mixed_image_button = gr.Button("generate_mixed_image", variant='primary')
+                clear_btn_image_to_image = gr.Button("Clear", variant="secondary")
+
+                gr.Markdown(description_generate_mixed_image)
+                # gr.Examples(examples = ["examples/12830823_87d2654e31.jpg", "examples/27782020_4dab210360.jpg", "examples/44129946_9eeb385d77.jpg"],
+                #             inputs=[text_input,red_color_slider_1,green_color_slider_1,blue_color_slider_1,noised_image_checkbox],
+                #             outputs=image_input4,
+                #             fn=run_generate_mixed_image,
+                #             cache_examples=True,
+                #             examples_per_page=4)
+
+    run_generate_mixed_image_button.click(run_generate_mixed_image,
+                        inputs=[text_input2,text_input3,noised_image_checkbox_2],
+                        outputs=image_input4)
+
+    ####################################################################################################################
+    with gr.Tab("progress_video"):
+        # Images
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                text_input5.render()
+
+            with gr.Column(scale=1):
+                video_output.render()
+
+        # Submit & Clear
+        with gr.Row():
+            with gr.Column():
+                run_progress_video_button = gr.Button("progress_video", variant='primary')
+                clear_btn_progress_video = gr.Button("Clear", variant="secondary")
+
+                gr.Markdown(description_progress_video)
+                # gr.Examples(examples = ["examples/12830823_87d2654e31.jpg", "examples/27782020_4dab210360.jpg", "examples/44129946_9eeb385d77.jpg"],
+                #             inputs=[text_input5],
+                #             outputs=video_output,
+                #             fn=run_process_video,
+                #             examples_per_page=4)
+
+    run_progress_video_button.click(run_process_video,
+                        inputs=[
+                            text_input5,
+                        ],
+                        outputs=video_output)
+    
+    #######################################################################################################################
+    #######################################################################################################################
+    def clear():
+        return None, None
+    
+    def clear_text():
+        return None, None, None
+
+    clear_btn_text_to_image.click(clear, outputs=[image_input1, image_input1])
+    clear_btn_image_to_image.click(clear, outputs=[image_input2, image_input3])
+    clear_btn_progress_video.click(clear, outputs=[image_input2, image_input3])
+demo.queue()
+demo.launch(debug=True)

main_inference.py

@@ -0,0 +1,252 @@
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer
+from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler
+from tqdm.auto import tqdm
+from torch import autocast
+from PIL import Image
+from matplotlib import pyplot as plt
+import numpy
+from torchvision import transforms as tfms
+import shutil
+# For video display:
+import cv2
+from IPython.display import HTML
+from base64 import b64encode
+import os 
+from utils import color_loss,latents_to_pil,pil_to_latent,sketch_loss
+# Set device
+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)
+
+scheduler.set_timesteps(15)
+
+def generate_mixed_image(prompt1, prompt2,noised_image=False):
+    mix_factor = 0.4 #@param
+    height = 512                        # default height of Stable Diffusion
+    width = 512                         # default width of Stable Diffusion
+    num_inference_steps = 50  #@param           # Number of denoising steps
+    guidance_scale = 8                # Scale for classifier-free guidance
+    generator = torch.manual_seed(32)   # Seed generator to create the inital latent noise
+    batch_size = 1
+
+    # Prep text
+    # Embed both prompts
+    text_input1 = tokenizer([prompt1], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        text_embeddings1 = text_encoder(text_input1.input_ids.to(torch_device))[0]
+        text_input2 = tokenizer([prompt2], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        text_embeddings2 = text_encoder(text_input2.input_ids.to(torch_device))[0]
+        # Take the average
+    text_embeddings = (text_embeddings1*mix_factor + \
+                    text_embeddings2*(1-mix_factor))
+    # And the uncond. input as before:
+    max_length = max(text_input1.input_ids.shape[-1],text_input2.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
+    scheduler.set_timesteps(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.sigmas[0] # Need to scale to match k
+
+    # Loop
+    with autocast("cuda"):
+        for i, t in tqdm(enumerate(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 = latent_model_input / ((sigma**2 + 1) ** 0.5)
+
+            # 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, i, latents)["prev_sample"]
+            
+    if noised_image:
+        output = generate_noised_version_of_image(latents_to_pil(latents,vae)[0])
+    else:
+        output = latents_to_pil(latents,vae)[0]
+
+    return output
+
+def generate_image(prompt,color_postprocessing=False,postporcessing_color=None,color_loss_scale=40,noised_image=False):
+    #@title Store the predicted outputs and next frame for later viewing
+    #prompt = 'A campfire (oil on canvas)' #
+    height = 512                        # default height of Stable Diffusion
+    width = 512                         # default width of Stable Diffusion
+    num_inference_steps = 50  #          # Number of denoising steps
+    guidance_scale = 8 #         # Scale for classifier-free guidance
+    generator = torch.manual_seed(32)   # Seed generator to create the inital latent noise
+    batch_size = 1
+    
+    # Define the directory name
+    directory_name = "steps"
+
+    # Check if the directory exists, and if so, delete it
+    if os.path.exists(directory_name):
+        shutil.rmtree(directory_name)
+    
+    #Create the directory
+    os.makedirs(directory_name)
+    # Prep text
+    text_input = tokenizer([prompt], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
+
+    # And the uncond. input as before:
+    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
+    scheduler.set_timesteps(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.sigmas[0] # Need to scale to match k
+
+    # Loop
+    with autocast("cuda"):
+        for i, t in tqdm(enumerate(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 = latent_model_input / ((sigma**2 + 1) ** 0.5)
+
+            # 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 ###
+            # Requires grad on the latents
+            if color_postprocessing:
+                latents = latents.detach().requires_grad_()
+
+                # Get the predicted x0:
+                latents_x0 = latents - sigma * noise_pred
+
+                # Decode to image space
+                denoised_images = vae.decode((1 / 0.18215) * latents_x0) / 2 + 0.5 # (0, 1)
+
+                # Calculate loss
+                #loss = sketch_loss(denoised_images) * color_loss_scale
+                loss = color_loss(denoised_images,postporcessing_color) * color_loss_scale
+                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
+
+
+                ### And saving as before ###
+                # Get the predicted x0:
+                latents_x0 = latents - sigma * noise_pred
+                im_t0 = latents_to_pil(latents_x0,vae)[0]
+
+                # And the previous noisy sample x_t -> x_t-1
+                latents = scheduler.step(noise_pred, i, latents)["prev_sample"]
+                im_next = latents_to_pil(latents,vae)[0]
+
+                # Combine the two images and save for later viewing
+                im = Image.new('RGB', (1024, 512))
+                im.paste(im_next, (0, 0))
+                im.paste(im_t0, (512, 0))
+                im.save(f'steps/{i:04}.jpeg')
+                
+            else:
+                latents = scheduler.step(noise_pred, i, latents)["prev_sample"]
+
+    if noised_image:
+        output = generate_noised_version_of_image(latents_to_pil(latents,vae)[0])
+    else:
+        output = latents_to_pil(latents,vae)[0]
+
+    return output
+
+def progress_video(prompt):
+    pil_image = generate_image(prompt)
+    # Generate a list of image file paths (replace with your own logic)
+    num_frames = len(os.listdir("steps/"))
+    image_files = [f"steps/{i:04d}.jpeg" for i in range(1, num_frames + 1)]
+    # Read the first image to get its size (assuming all images have the same size)
+    first_image = cv2.imread({image_files[0]})
+    height, width, _ = first_image.shape
+
+    # Define the output video writer
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for MP4
+    out = cv2.VideoWriter('out.mp4', fourcc, 12, (width, height))
+
+    for image_file in image_files:
+        frame = cv2.imread(image_file)
+        out.write(frame)
+
+    out.release()
+    return "out.mp4"
+
+def generate_noised_version_of_image(pil_image):
+    # View a noised version
+    encoded = pil_to_latent(pil_image,vae)
+    noise = torch.randn_like(encoded) # Random noise
+    timestep = 150 # i.e. equivalent to that at 150/1000 training steps
+    encoded_and_noised = scheduler.add_noise(encoded, noise, timestep)
+    return latents_to_pil(encoded_and_noised,vae)[0] # Display
+
+
+# if __name__ == "__main__":
+#     prompt = 'A campfire (oil on canvas)'
+#     color_loss_scale = 40
+#     color_postprocessing = False
+#     pil_image = generate_mixed_image("a dog", "a cat")
+#     #pil_image = generate_image(prompt,color_postprocessing,color_loss_scale)
+#     #pil_image = generate_noised_version_of_image(Image.open('output.png').resize((512, 512)))
+#     pil_image.save("output1.png")
+    
+if __name__ == "__main__":
+    progress_video("lol")

requirements.txt

@@ -0,0 +1,8 @@
+transformers
+diffusers==0.2.4 
+sentence_transformers
+gradio
+torch
+torchvision
+matplotlib
+opencv-python

utils.py

@@ -0,0 +1,75 @@
+from torchvision.transforms import functional as F
+from torchvision import transforms as tfms
+from PIL import Image, ImageEnhance
+#from legofy import legofy_image
+import numpy as np
+from torchvision.transforms import functional as F
+from PIL import Image, ImageEnhance
+import torch 
+import cv2 
+
+to_tensor_tfm = tfms.ToTensor()
+torch_device = "cpu"
+
+
+def pil_to_latent(input_im,vae):
+    
+  # Single image -> single latent in a batch (so size 1, 4, 64, 64)
+  with torch.no_grad():
+    latent = vae.encode(to_tensor_tfm(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling
+  return 0.18215 * latent.mode() # or .mean or .sample
+
+def latents_to_pil(latents,vae):
+  # bath of latents -> list of images
+  latents = (1 / 0.18215) * latents
+  with torch.no_grad():
+    image = vae.decode(latents)
+  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 color_loss(images,color):
+  # Scale the coming color 
+  red,green,blue = (color[0]/255)*0.9,(color[1]/255)*0.9,(color[2]/255)*0.9
+  
+  red_chennel_error = torch.abs(images[:,0, :, :] - red).mean()
+  green_chennel_error = torch.abs(images[:,1, :, :] - green).mean()
+  blue_chennel_error = torch.abs(images[:,2, :, :] - blue).mean()
+  print(red_chennel_error, green_chennel_error, blue_chennel_error)
+  error = red_chennel_error + green_chennel_error + blue_chennel_error
+  return error
+
+import torch
+from PIL import Image, ImageOps, ImageFilter
+import torchvision.transforms as transforms
+
+def sketch_loss(image):
+    # Convert PyTorch tensor to a PIL image
+    to_pil = transforms.ToPILImage()
+    pil_image = to_pil(image[0])
+
+    # Convert the PIL image to grayscale
+    gray_image = ImageOps.grayscale(pil_image)
+
+    # Apply an inverted pencil sketch effect
+    inverted_image = ImageOps.invert(gray_image)
+
+    # Apply a blur effect to smooth the sketch
+    pencil_sketch = inverted_image.filter(ImageFilter.GaussianBlur(radius=5))
+
+    # Convert the PIL image back to a PyTorch tensor
+    to_tensor = transforms.ToTensor()
+    sketch_tensor = to_tensor(pencil_sketch).unsqueeze(0)
+    sketch_tensor.requires_grad = True  # Enable gradients
+
+    #if num_channels == 3:
+    #    # If the input was originally in CHW format (3 channels), permute it to CHW
+    sketch_tensor = sketch_tensor.permute(0, 3, 1, 2)
+
+    # Calculate the loss based on the watercolour_image tensor
+    loss = torch.abs(sketch_tensor - 0.9).mean()  # Modify 0.5 to your desired threshold
+
+    return loss