Update app.py

app.py

@@ -115,27 +115,36 @@ def generate_image(seed, prompt, loss_type, loss_flag=False):
 
     return latents
 
-# Gradio interface function
-def generate_images(prompt, style, guidance_type):
-    images = show_image(prompt, style, guidance_type)
-    return images[0], images[1]
+def generate_image(prompt, style, guidance_type):
+    styled_prompt = f"{prompt} in the style of {styles_mapping[style]}"
+    seed = torch.randint(0, 1000000, (1,)).item()
+    latents = generate_image(seed, styled_prompt, guidance_type, loss_flag=True)
+    with torch.no_grad():
+        image = sd_pipeline.decode_latents(latents)
+    image = sd_pipeline.numpy_to_pil(image)[0]
+    return image
+
+def get_examples():
+    examples = [
+        ["A bird sitting on a tree", "Midjourney", "edge"],
+        ["Cats fighting on the road", "Marc Allante", "brightness"],
+        ["A mouse with the head of a puppy", "Hitokomoru Style", "contrast"],
+        ["A woman with a smiling face in front of an Italian Pizza", "Hanfu Anime", "brightness"],
+        ["A campfire (oil on canvas)", "Birb Style", "blue"],
+    ]
+    return examples
 
-# Create Gradio interface
 iface = gr.Interface(
-    fn=generate_images,
+    fn=generate_image,
     inputs=[
         gr.Textbox(label="Prompt"),
         gr.Dropdown(list(styles_mapping.keys()), label="Style"),
         gr.Dropdown(["blue", "edge", "contrast", "brightness", "sharpness", "saturation"], label="Guidance Type"),
     ],
-    outputs=[
-        gr.Image(label="Image without Loss"),
-        gr.Image(label="Image with Loss"),
-    ],
+    outputs=gr.Image(label="Generated Image"),
+    title="Stable Diffusion with Custom Styles",
+    description="Generate images using a custom Stable Diffusion model with various styles and guidance types.",
     examples=get_examples(),
-    title="Text Inversion Image Generation",
-    description="Generate images using text inversion with different styles and guidance types.",
 )
 
-# Launch the app
 iface.launch()