code update

app.py

@@ -42,11 +42,13 @@ style_files = ['stable_diffusion/learned_embeddings/arcane-style-jv.bin', 'stabl
                      'stable_diffusion/learned_embeddings/dr-strange.bin', 'stable_diffusion/learned_embeddings/midjourney-style.bin',
                      'stable_diffusion/learned_embeddings/oil_style.bin']
 
+images_without_loss = []
+images_with_loss = []
 
 seed_values = [8,16,50,80,128]
 height = 512                        # default height of Stable Diffusion
 width = 512                         # default width of Stable Diffusion
-num_inference_steps = 1           # Number of denoising steps
+num_inference_steps = 10            # Number of denoising steps
 guidance_scale = 7.5                # Scale for classifier-free guidance
 num_styles = len(style_files)
 
@@ -80,7 +82,6 @@ def get_output_embeds(input_embeddings):
     # And now they're ready!
     return output
 
-
 def get_style_embeddings(style_file):
     style_embed = torch.load(style_file)
     style_name = list(style_embed.keys())[0]
@@ -247,55 +248,48 @@ def display_images_in_rows(images_with_titles, titles):
                 axes[r, c].axis('off')
 
     return fig
+    # plt.show()
+
 
 def image_generator(prompt = "dog", loss_function=None):
-    
-    if loss_function == "Yes":
-      loss_function = vibrance_loss
-    else:
-      loss_function = None
-        
-    images_without_loss = []
-    images_with_loss = []
-    
-    seed_values = [8,16,50,80,128]
-    height = 512                        # default height of Stable Diffusion
-    width = 512                         # default width of Stable Diffusion
-    num_inference_steps = 1           # Number of denoising steps
-    guidance_scale = 7.5                # Scale for classifier-free guidance
-    num_styles = len(style_files)
-    
-    for i in range(num_styles):
-      this_generated_img_1 = generate_images(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = None)
-      images_without_loss.append(this_generated_img_1)
+  images_without_loss = []
+  images_with_loss = []
+  if loss_function == "Yes":
+    loss_function = vibrance_loss
+  else:
+    loss_function = None
+
+  for i in range(num_styles):
+      generated_img = generate_images(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = None)
+      images_without_loss.append(generated_img)
       if loss_function:
-        this_generated_img_2 = generate_images(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = loss_function)
-        images_with_loss.append(this_generated_img_2)
-    
-    generated_sd_images = []
-    titles = ["Arcane Style", "Birb Style", "Dr Strange Style", "Midjourney Style", "Oil Style"]
-    
-    for i in range(len(titles)):
-        generated_sd_images.append((images_without_loss[i], titles[i])) 
-        if images_with_loss != []:
-          generated_sd_images.append((images_with_loss[i], titles[i])) 
-            
-    return display_images_in_rows(generated_sd_images, titles)
+        generated_img = generate_images(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = loss_function)
+        images_with_loss.append(generated_img)
+
+  generated_sd_images = []
+  titles = ["Arcane Style", "Birb Style", "Dr Strange Style", "Midjourney Style", "Oil Style"]
+
+  for i in range(len(titles)):
+    generated_sd_images.append((images_without_loss[i], titles[i])) 
+    if images_with_loss != []:
+      generated_sd_images.append((images_with_loss[i], titles[i])) 
+
 
+  return display_images_in_rows(generated_sd_images, titles)
 
 # Create a wrapper function for show_misclassified_images()
 def image_generator_wrapper(prompt = "dog", loss_function=None):
-    if loss_function == "Yes":
-      loss_function = vibrance_loss
-    else:
-      loss_function = None
+  if loss_function == "Yes":
+    loss_function = vibrance_loss
+  else:
+    loss_function = None
 
-    return image_generator(prompt, loss_function)
+  return image_generator(prompt, loss_function)
 
-description = "Generate an image with a prompt and apply loss if you wish to"
+description = "Generate an image with a prompt and apply vibrance loss if you wish to"
 
 demo = gr.Interface(image_generator,
-                    inputs=[gr.Textbox(label="Enter prompt for generating", type="text", value="dog sitting on a bench"),
+                    inputs=[gr.Textbox(label="Enter prompt for generation", type="text", value="dog sitting on a bench"),
                             gr.Radio(["Yes", "No"], value="No"  , label="Apply vibrance loss")],
-                    outputs=gr.Plot(label="Outputs"), title = "Stable Diffusion", description=description)
-demo.launch()
+                    outputs=gr.Plot(label="Generated Images"), title = "Stable Diffusion", description=description)
+demo.launch()