Spaces:
Runtime error
Runtime error
File size: 4,452 Bytes
d0a7bb3 |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import torch
import gradio as gr
from functools import partial
from diffusers_patch import OMSPipeline
def create_sdxl_lcm_lora_pipe(sd_pipe_name_or_path, oms_name_or_path, lora_name_or_path):
from diffusers import StableDiffusionXLPipeline, LCMScheduler
sd_pipe = StableDiffusionXLPipeline.from_pretrained(sd_pipe_name_or_path, torch_dtype=torch.float16, variant="fp16", add_watermarker=False).to('cuda')
print('successfully load pipe')
sd_scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config)
sd_pipe.load_lora_weights(lora_name_or_path, variant="fp16")
pipe = OMSPipeline.from_pretrained(oms_name_or_path, sd_pipeline = sd_pipe, torch_dtype=torch.float16, variant="fp16", trust_remote_code=True, sd_scheduler=sd_scheduler)
pipe.to('cuda')
return pipe
class GradioDemo:
def __init__(
self,
sd_pipe_name_or_path = "stabilityai/stable-diffusion-xl-base-1.0",
oms_name_or_path = 'h1t/oms_b_openclip_xl',
lora_name_or_path = 'latent-consistency/lcm-lora-sdxl'
):
self.pipe = create_sdxl_lcm_lora_pipe(sd_pipe_name_or_path, oms_name_or_path, lora_name_or_path)
def _inference(
self,
prompt = None,
oms_prompt = None,
oms_guidance_scale = 1.0,
num_inference_steps = 4,
sd_pipe_guidance_scale = 1.0,
seed = 1024,
):
pipe_kwargs = dict(
prompt = prompt,
num_inference_steps = num_inference_steps,
guidance_scale = sd_pipe_guidance_scale,
)
generator = torch.Generator(device=self.pipe.device).manual_seed(seed)
pipe_kwargs.update(oms_flag=False)
print(f'raw kwargs: {pipe_kwargs}')
image_raw = self.pipe(
**pipe_kwargs,
generator=generator
)['images'][0]
generator = torch.Generator(device=self.pipe.device).manual_seed(seed)
pipe_kwargs.update(oms_flag=True, oms_prompt=oms_prompt, oms_guidance_scale=1.0)
print(f'w/ oms wo/ cfg kwargs: {pipe_kwargs}')
image_oms = self.pipe(
**pipe_kwargs,
generator=generator
)['images'][0]
oms_guidance_flag = oms_guidance_scale != 1.0
if oms_guidance_flag:
generator = torch.Generator(device=self.pipe.device).manual_seed(seed)
pipe_kwargs.update(oms_guidance_scale=oms_guidance_scale)
print(f'w/ oms +cfg kwargs: {pipe_kwargs}')
image_oms_cfg = self.pipe(
**pipe_kwargs,
generator=generator
)['images'][0]
else:
image_oms_cfg = None
return image_raw, image_oms, image_oms_cfg, gr.update(visible=oms_guidance_flag)
def mainloop(self):
with gr.Blocks() as demo:
gr.Markdown("# One More Step Demo")
with gr.Row():
with gr.Column():
prompt = gr.Textbox(label="Prompt", value="a cat")
oms_prompt = gr.Textbox(label="OMS Prompt", value="orange car")
oms_guidance_scale = gr.Slider(label="OMS Guidance Scale", minimum=1.0, maximum=5.0, value=1.5, step=0.1)
run_button = gr.Button(value="Generate images")
with gr.Accordion("Advanced options", open=False):
num_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=4, step=1)
sd_guidance_scale = gr.Slider(label="SD Pipe Guidance Scale", minimum=0.1, maximum=30.0, value=1.0, step=0.1)
seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=False, value=1024)
with gr.Column():
output_raw = gr.Image(label="SDXL w/ LCM-LoRA w/o OMS ")
output_oms = gr.Image(label="w/ OMS w/o OMS CFG")
with gr.Column(visible=False) as oms_cfg_wd:
output_oms_cfg = gr.Image(label=f"w/ OMS w/ OMS CFG")
ips = [prompt, oms_prompt, oms_guidance_scale, num_steps, sd_guidance_scale, seed]
run_button.click(fn=self._inference, inputs=ips, outputs=[output_raw, output_oms, output_oms_cfg, oms_cfg_wd])
demo.queue(max_size=20)
demo.launch()
if __name__ == "__main__":
gradio_demo = GradioDemo()
gradio_demo.mainloop()
|