Update app.py

app.py

@@ -2,15 +2,10 @@ import gradio as gr
 from diffusers import StableDiffusionXLPipeline, DDIMScheduler
 import torch
 import numpy as np
-from diffusers.utils import load_image
 from PIL import Image
 import io
 import sys
 import os
-current_dir = os.path.dirname(os.path.realpath(__file__))
-if current_dir not in sys.path:
-    sys.path.append(current_dir)
-
 import sa_handler
 import inversion
 
@@ -26,10 +21,10 @@ pipeline = StableDiffusionXLPipeline.from_pretrained(
 ).to("cuda")
 
 # Function to process the image
-def process_image(image, prompt, style):
-    src_prompt = f'Man laying in a bed, {style}.'
-    
-    num_inference_steps = 50
+def process_image(image, prompt, style, src_description, inference_steps, shared_score_shift, shared_score_scale, guidance_scale):
+    src_prompt = f'{src_description}, {style}.'
+
+    num_inference_steps = inference_steps
     x0 = np.array(Image.fromarray(image).resize((1024, 1024)))
     zts = inversion.ddim_inversion(pipeline, x0, src_prompt, num_inference_steps, 2)
 
@@ -38,9 +33,6 @@ def process_image(image, prompt, style):
         f"{prompt}, {style}."
     ]
 
-    shared_score_shift = np.log(2)
-    shared_score_scale = 1.0
-
     handler = sa_handler.Handler(pipeline)
     sa_args = sa_handler.StyleAlignedArgs(
         share_group_norm=True, share_layer_norm=True, share_attention=True,
@@ -59,23 +51,27 @@ def process_image(image, prompt, style):
 
     images_a = pipeline(prompts, latents=latents,
                         callback_on_step_end=inversion_callback,
-                        num_inference_steps=num_inference_steps, guidance_scale=10.0).images
+                        num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images
 
     handler.remove()
 
     return Image.fromarray(images_a[1])
 
-# Gradio interface
 iface = gr.Interface(
     fn=process_image,
     inputs=[
-        gr.inputs.Image(type="numpy"), 
-        gr.inputs.Textbox(label="Enter your prompt"),
-        gr.inputs.Textbox(label="Enter your style", default="medieval painting")
+        gr.Image(type="numpy"), 
+        gr.Textbox(label="Enter your prompt"),
+        gr.Textbox(label="Enter your style", default="medieval painting"),
+        gr.Textbox(label="Enter source description", default="Man laying in a bed"),
+        gr.Slider(minimum=5, maximum=50, step=1, default=50, label="Number of Inference Steps"),
+        gr.Slider(minimum=1, maximum=2, step=0.01, default=1.5, label="Shared Score Shift"),
+        gr.Slider(minimum=0, maximum=1, step=0.01, default=0.5, label="Shared Score Scale"),
+        gr.Slider(minimum=5, maximum=120, step=1, default=10, label="Guidance Scale")
     ],
     outputs="image",
     title="Stable Diffusion XL with Style Alignment",
     description="Generate images in the style of your choice."
 )
 
-iface.launch()
+iface.launch()