Create app.py

app.py

@@ -0,0 +1,75 @@
+import gradio as gr
+from diffusers import StableDiffusionXLPipeline, DDIMScheduler
+import torch
+import sa_handler
+import inversion
+import numpy as np
+from diffusers.utils import load_image
+from PIL import Image
+import io
+
+# 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_prompt = f'Man laying in a bed, {style}.'
+    
+    num_inference_steps = 50
+    x0 = np.array(Image.fromarray(image).resize((1024, 1024)))
+    zts = inversion.ddim_inversion(pipeline, x0, src_prompt, num_inference_steps, 2)
+
+    prompts = [
+        src_prompt,
+        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,
+        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)
+
+    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')
+    latents[0] = zT
+
+    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()
+
+    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")
+    ],
+    outputs="image",
+    title="Stable Diffusion XL with Style Alignment",
+    description="Generate images in the style of your choice."
+)
+
+iface.launch()