main app.py file.

app.py

@@ -0,0 +1,65 @@
+import os
+import sys
+import sys
+import torch
+import gradio as gr
+from stable_diffusion import StableDiffusion
+from utils import invert_loss, get_style_embeddings, show_images
+
+torch.manual_seed(1)
+torch_device = "cuda" if torch.cuda.is_available() else "cpu" 
+
+
+def gradio_interface(prompt:str, style_name:str):
+    stable_diffuser = StableDiffusion(torch_device)
+    #outputs_1 = []
+    #outputs_2 = []
+    seed_values = [1,2,3,4,5]
+    custom_loss_scale = 150.0
+    num_styles = len(style_files)
+    style_file = style_files[style_name]
+    style_token_embedding = get_style_embeddings(style_file) if style_file is not None else None
+    this_generated_img_1 = stable_diffuser.generate_image_with_custom_style(prompt,
+                                                          style_token_embedding = style_token_embedding,
+                                                          random_seed = 42,
+                                                          custom_loss_fn = None)
+    #outputs_1.append(this_generated_img_1)
+    this_generated_img_2 = stable_diffuser.generate_image_with_custom_style(prompt,
+                                                          style_token_embedding = style_token_embedding,
+                                                          random_seed = 42,
+                                                          custom_loss_fn = invert_loss,
+                                                          custom_loss_scale = custom_loss_scale)
+    #outputs_2.append(this_generated_img_2)
+    return this_generated_img_1, this_generated_img_2
+    
+
+style_files = {'style_1': None,
+               'watercolor': 'learned_embeds_watercolor.bin',
+               'strip': 'learned_embeds_strip_style.bin',
+               'oil_paint': 'learned_embeds_oil_paint.bin',
+               'kaleido': 'learned_embeds_kaleido.bin',
+                'doodle': 'learned_embeds_doodle.bin'}
+
+torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+prompt_examples = [["A cat wearing a party hat vivid colors", 'style_1']]
+#["A dog wearing sunglasses on a skateboard"],
+#["An oil painting of a bear playing guitar"]]
+#["A tiger dressed in a hip hop styled clothing"],
+#["An oil painting of a lion eating pizza"]]
+
+# Define Interface
+description = 'A Stable Diffusion based Generative AI tool to generate images. Also generates grayscale images.'
+title = 'Image Generation using Stable Diffusion with styles.'
+style_names = ['No_style','watercolor','strip','oil_paint','kaleido','doodle']
+demo = gr.Interface(gradio_interface,
+                    inputs = [gr.Textbox('A deer crossing the street', label="Text Prompt"),
+                             gr.Dropdown(style_names,value='style_1',label="Display Style")],
+                    outputs = [gr.Image(shape=(512, 512),label='Generated Image').style(width=512, height=512),
+                               gr.Image(shape=(512, 512),label='Generated Images with custom loss').style(width=512, height=512)
+                              ],
+                    examples=prompt_examples,
+                    title = title,
+                    description = description
+                   )
+demo.launch(debug=True)