Update app.py

app.py

@@ -1,38 +1,24 @@
-import gradio as gr
 from diffusers import StableDiffusionXLPipeline, DDIMScheduler
 import torch
+import gradio as gr
+import inversion
 import numpy as np
 from PIL import Image
-import io
-import sys
-import os
 import sa_handler
-import inversion
-
-# Model Load
-scheduler = DDIMScheduler(
-    beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
-    clip_sample=False, set_alpha_to_one=False)
-
-pipeline = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16",
-    use_safetensors=True,
-    scheduler=scheduler
-).to("cuda")
 
-# Function to process the image
-def process_image(image, prompt, style, src_description, inference_steps, shared_score_shift, shared_score_scale, guidance_scale):
-    src_prompt = f'{src_description}, {style}.'
+device = "cuda" if torch.cuda.is_available() else "cpu"
+scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
+pipeline = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True, scheduler=scheduler).to(device)
 
-    num_inference_steps = inference_steps
-    x0 = np.array(Image.fromarray(image).resize((1024, 1024)))
+def run(image, src_style, src_prompt, prompts, shared_score_shift, shared_score_scale, guidance_scale, num_inference_steps, large, seed):
+    prompts = prompts.splitlines()
+    dim, d = (1024, 128) if large else (512, 64)
+    image = image.resize((dim, dim))
+    x0 = np.array(image)
     zts = inversion.ddim_inversion(pipeline, x0, src_prompt, num_inference_steps, 2)
+    prompts.insert(0, src_prompt)
 
-    prompts = [
-        src_prompt,
-        f"{prompt}, {style}."
-    ]
-
+    shared_score_shift = np.log(shared_score_shift)
     handler = sa_handler.Handler(pipeline)
     sa_args = sa_handler.StyleAlignedArgs(
         share_group_norm=True, share_layer_norm=True, share_attention=True,
@@ -40,38 +26,43 @@ def process_image(image, prompt, style, src_description, inference_steps, shared
         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)
+    for i in range(1, len(prompts)):
+        prompts[i] = f'{prompts[i]}, {src_style}.'
 
+    zT, inversion_callback = inversion.make_inversion_callback(zts, offset=5)
     g_cpu = torch.Generator(device='cpu')
-    g_cpu.manual_seed(10)
-
-    latents = torch.randn(len(prompts), 4, 128, 128, device='cpu', generator=g_cpu,
-                          dtype=pipeline.unet.dtype,).to('cuda:0')
+    if seed > 0:
+        g_cpu.manual_seed(seed)
+    latents = torch.randn(len(prompts), 4, d, d, device='cpu', generator=g_cpu, dtype=pipeline.unet.dtype,).to(device)
     latents[0] = zT
 
-    images_a = pipeline(prompts, latents=latents,
-                        callback_on_step_end=inversion_callback,
-                        num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images
+    images_a = pipeline(prompts, latents=latents, callback_on_step_end=inversion_callback, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images
 
     handler.remove()
+    torch.cuda.empty_cache()
+    images_pil = [Image.fromarray((img * 255).astype(np.uint8)) for img in images_a]
+    return images_pil
 
-    return Image.fromarray(images_a[1])
-
-iface = gr.Interface(
-    fn=process_image,
-    inputs=[
-        gr.Image(type="numpy"), 
-        gr.Textbox(label="Enter your prompt"),
-        gr.Textbox(label="Enter your style", value="medieval painting"),
-        gr.Textbox(label="Enter source description", value="Man laying in a bed"),
-        gr.Slider(minimum=5, maximum=50, step=1, value=50, label="Number of Inference Steps"),
-        gr.Slider(minimum=1, maximum=2, step=0.01, value=1.5, label="Shared Score Shift"),
-        gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label="Shared Score Scale"),
-        gr.Slider(minimum=5, maximum=120, step=1, value=10, label="Guidance Scale")
-    ],
-    outputs="image",
-    title="Stable Diffusion XL with Style Alignment",
-    description="Generate images in the style of your choice."
-)
+with gr.Blocks() as demo:
+    with gr.Markdown("""
+    # Welcome to Tonic's Stable Style Align
+    Add a reference picture , describe the style and add prompts to generate images in that style. It's the most interesting with your own art ! 
+    """)
+    with gr.Row():
+        gr.Image(label="Reference image", type="pil")
+    with gr.Row():
+        gr.Textbox(label="Describe the reference style")
+        gr.Textbox(label="Describe the reference image")
+        gr.Textbox(label="Prompts to generate images (separate with new lines)", lines=5)
+    with gr.Accordion(label="Advanced Settings"):
+        with gr.Row():
+            gr.Number(value=1.1, label="shared_score_shift", min=1.0, max=2.0)
+            gr.Number(value=1.0, label="shared_score_scale", min=0.0, max=1.0)
+            gr.Number(value=10.0, label="guidance_scale", min=5.0, max=20.0)
+            gr.Number(value=50, label="num_inference_steps", min=1, max=100, precision=0)
+            gr.Checkbox(False, label="Large (1024x1024)")
+            gr.Number(value=0, label="seed (0 for random)", min=0, max=10000, precision=0)
+    with gr.Row():
+        gr.Gallery()
 
-iface.launch()
+demo.launch()