Upload 2 files

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from PIL import Image
 import random
 
-from utils import load_learned_embeds, style_guide
+from utils import load_learned_embeds, style_guide, IMAGE_SIZE
 
 def generate_output_image(query, selection):
     """Simulates an output image generation based on query and selection."""
@@ -25,7 +25,7 @@ def check_visibility(query, selection):
 
 def ui():
     with gr.Blocks() as demo:
-        gr.Markdown("## Image Selection & Query-Based Search")
+        gr.Markdown("## Image Generation using Stable Diffusion")
 
         # First Block: Displaying Images with Labels
 
@@ -36,18 +36,21 @@ def ui():
             for i in range(6):
                 image = f'./examples/{i}.png'
                 if i == 0:
-                    label = 'Image Sharpness Loss'
+                    label = 'Sharpness Loss (Additional Guidance)'
                 else:
                     label = style_guide[i-1][0]
 
                 with gr.Column():
-                    gr.Image(value=image, label=f"{label}")
+                    gr.Image(value=image, label=f"{label}",
+                             height=IMAGE_SIZE,
+                             show_download_button=False,
+                             show_share_button=False)
                     #gr.Textbox(value=f"Label {label}", interactive=False)
 
         # Second Block: Query Input and Dropdown
         with gr.Row():
             query_input = gr.Textbox(label="Enter Query")
-            dropdown = gr.Dropdown(['None', "Sharpness Loss", 'Oil Style', 'Matrix Style', 'Stripe Style','Dreamy Painting Style','Polygon HD Style'],
+            dropdown = gr.Dropdown(['None', "Sharpness Loss (Additional Guidance)", 'Oil Style', 'Matrix Style', 'Stripe Style','Dreamy Painting Style','Polygon HD Style'],
                                    label="Select Option")
 
         # Dynamic Search Button (Initially Hidden)

utils.py

@@ -7,6 +7,7 @@ from tqdm.auto import tqdm
 from transformers import CLIPTextModel, CLIPTokenizer, logging
 import os
 import torch.nn.functional as F
+import random
 
 style_guide = [
     ('Oil Style', 'oilstyle_learned_embeds.bin'),
@@ -16,6 +17,7 @@ style_guide = [
     ('Polygon HD Style', 'lowpolyhd_learned_embeds.bin')
 ]
 
+IMAGE_SIZE = 224
 
 # Prep Scheduler
 def set_timesteps(scheduler, num_inference_steps):
@@ -78,11 +80,11 @@ def get_output_embeds(input_embeddings):
 
 
 def generate_with_embs(text_embeddings, add_guidance_loss=None, add_guidance_loss_scale=200):
-    height = 512  # default height of Stable Diffusion
-    width = 512  # default width of Stable Diffusion
+    height = IMAGE_SIZE  # default height of Stable Diffusion
+    width = IMAGE_SIZE  # default width of Stable Diffusion
     num_inference_steps = 30  # Number of denoising steps
     guidance_scale = 7.5  # Scale for classifier-free guidance
-    generator = torch.manual_seed(32)  # Seed generator to create the inital latent noise
+    generator = torch.manual_seed(random.randint(1, 10000))  # Seed generator to create the inital latent noise
     batch_size = 1
 
     max_length = text_embeddings.shape[1]