File size: 2,117 Bytes
e077396 c8497e9 e077396 |
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 |
import argparse
import time
import gradio as gr
import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel
from scheduling_dmd import DMDScheduler
parser = argparse.ArgumentParser()
parser.add_argument("--unet-path", type='Lykon/dreamshaper-8')
parser.add_argument("--model-path", type='aaronb/dreamshaper-8-dmd-kl-only-6kstep')
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
unet = UNet2DConditionModel.from_pretrained(args.unet_path)
pipe = DiffusionPipeline.from_pretrained(args.model_path, unet=unet)
pipe.scheduler = DMDScheduler.from_config(pipe.scheduler.config)
pipe.to(device=device, dtype=torch.float16)
def predict(prompt, seed=1231231):
generator = torch.manual_seed(seed)
last_time = time.time()
image = pipe(
prompt,
num_inference_steps=1,
guidance_scale=0.0,
generator=generator,
).images[0]
print(f"Pipe took {time.time() - last_time} seconds")
return image
css = """
#container{
margin: 0 auto;
max-width: 40rem;
}
#intro{
max-width: 100%;
text-align: center;
margin: 0 auto;
}
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="container"):
gr.Markdown(
"""# Distribution Matching Distillation
""",
elem_id="intro",
)
with gr.Row():
with gr.Row():
prompt = gr.Textbox(placeholder="Insert your prompt here:", scale=5, container=False)
generate_bt = gr.Button("Generate", scale=1)
image = gr.Image(type="filepath")
with gr.Accordion("Advanced options", open=False):
seed = gr.Slider(randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1)
inputs = [prompt, seed]
generate_bt.click(fn=predict, inputs=inputs, outputs=image, show_progress=False)
prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
demo.queue(api_open=False)
demo.launch(show_api=False)
|