Spaces:
Paused
Paused
File size: 2,823 Bytes
3d446a2 cd2da70 3d446a2 53ad954 3d446a2 53ad954 3d446a2 53ad954 3d446a2 53ad954 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import gradio as gr
from diffusers import StableDiffusionXLPipeline, DDIMScheduler
import torch
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}.'
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)
prompts = [
src_prompt,
f"{prompt}, {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)
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=guidance_scale).images
handler.remove()
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", 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() |