Update app.py

app.py

@@ -1,6 +1,4 @@
 import gradio as gr
-import spaces
-
 from diffusers import StableDiffusionXLPipeline, DDIMScheduler
 import torch
 import sa_handler
@@ -27,69 +25,71 @@ from diffusers.utils import load_image
 import inversion
 import numpy as np
 
-src_style = "medieval painting"
-src_prompt = f'Man laying in a bed, {src_style}.'
-image_path = './example_image/medieval-bed.jpeg'
+def run(ref_path, ref_style, ref_prompt, prompt1, prompt2, prompt3):
+
+    src_style = f"{ref_style}"
+    src_prompt = f"{ref_prompt}, {src_style}."
+    image_path = f"{ref_path}"
 
-num_inference_steps = 50
-x0 = np.array(load_image(image_path).resize((1024, 1024)))
-zts = inversion.ddim_inversion(pipeline, x0, src_prompt, num_inference_steps, 2)
-#mediapy.show_image(x0, title="innput reference image", height=256)
+    num_inference_steps = 50
+    x0 = np.array(load_image(image_path).resize((1024, 1024)))
+    zts = inversion.ddim_inversion(pipeline, x0, src_prompt, num_inference_steps, 2)
+    #mediapy.show_image(x0, title="innput reference image", height=256)
 
-# run StyleAligned
-prompts = [
-    src_prompt,
-    "A man working on a laptop",
-    "A man eats pizza",
-    "A woman playig on saxophone",
-]
+    # run StyleAligned
+    prompts = [
+        src_prompt,
+        prompt1,
+        prompt2.
+        prompt3
+    ]
 
-# some parameters you can adjust to control fidelity to reference
-shared_score_shift = np.log(2)  # higher value induces higher fidelity, set 0 for no shift
-shared_score_scale = 1.0  # higher value induces higher, set 1 for no rescale
+    # some parameters you can adjust to control fidelity to reference
+    shared_score_shift = np.log(2)  # higher value induces higher fidelity, set 0 for no shift
+    shared_score_scale = 1.0  # higher value induces higher, set 1 for no rescale
 
-# for very famouse images consider supressing attention to refference, here is a configuration example:
-# shared_score_shift = np.log(1)
-# shared_score_scale = 0.5
+    # for very famouse images consider supressing attention to refference, here is a configuration example:
+    # shared_score_shift = np.log(1)
+    # shared_score_scale = 0.5
 
-for i in range(1, len(prompts)):
-    prompts[i] = f'{prompts[i]}, {src_style}.'
+    for i in range(1, len(prompts)):
+        prompts[i] = f'{prompts[i]}, {src_style}.'
 
-handler = sa_handler.Handler(pipeline)
-sa_args = sa_handler.StyleAlignedArgs(
-    share_group_norm=True, share_layer_norm=True, share_attention=True,
-    adain_queries=True, adain_keys=True, adain_values=False,
-    shared_score_shift=shared_score_shift, shared_score_scale=shared_score_scale,)
-handler.register(sa_args)
+    handler = sa_handler.Handler(pipeline)
+    sa_args = sa_handler.StyleAlignedArgs(
+        share_group_norm=True, share_layer_norm=True, share_attention=True,
+        adain_queries=True, adain_keys=True, adain_values=False,
+        shared_score_shift=shared_score_shift, shared_score_scale=shared_score_scale,)
+    handler.register(sa_args)
 
-zT, inversion_callback = inversion.make_inversion_callback(zts, offset=5)
+    zT, inversion_callback = inversion.make_inversion_callback(zts, offset=5)
 
-g_cpu = torch.Generator(device='cpu')
-g_cpu.manual_seed(10)
+    g_cpu = torch.Generator(device='cpu')
+    g_cpu.manual_seed(10)
 
-latents = torch.randn(len(prompts), 4, 128, 128, device='cpu', generator=g_cpu,
+    latents = torch.randn(len(prompts), 4, 128, 128, device='cpu', generator=g_cpu,
                       dtype=pipeline.unet.dtype,).to('cuda:0')
-latents[0] = zT
+    latents[0] = zT
 
-images_a = pipeline(prompts, latents=latents,
+    images_a = pipeline(prompts, latents=latents,
                     callback_on_step_end=inversion_callback,
                     num_inference_steps=num_inference_steps, guidance_scale=10.0).images
 
-handler.remove()
-mediapy.show_images(images_a, titles=[p[:-(len(src_style) + 3)] for p in prompts])
-
-
-@spaces.GPU
-def infer(prompts):
-
+    handler.remove()
+    #mediapy.show_images(images_a, titles=[p[:-(len(src_style) + 3)] for p in prompts])
 
-    images = pipeline(sets_of_prompts,).images
     return images_a
+    
 
 gr.Interface(
-    fn=infer,
+    fn=run,
     inputs=[
-        gr.Textbox(value="Hit submit button to test")
+        gr.Image(type="filepath", value="./example_image/medieval-bed.jpeg"),
+        gr.Textbox(value="medieval painting"),
+        gr.Textbox(value="Man laying on bed"),
+        gr.Textbox(value="A man working on a laptop"),
+        gr.Textbox(value="A man eating pizza"),
+        gr.Textbox(value="A woman playing on saxophone")
     ],
     outputs=[
         gr.Gallery()