|
import gradio as gr |
|
import torch |
|
from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler |
|
from huggingface_hub import hf_hub_download |
|
import spaces |
|
from PIL import Image |
|
|
|
|
|
base = "stabilityai/stable-diffusion-xl-base-1.0" |
|
repo = "tianweiy/DMD2" |
|
checkpoints = { |
|
"1-Step" : ["dmd2_sdxl_1step_unet_fp16.bin", 1], |
|
"4-Step" : ["dmd2_sdxl_4step_unet_fp16.bin", 4], |
|
} |
|
loaded = None |
|
|
|
CSS = """ |
|
.gradio-container { |
|
max-width: 690px !important; |
|
} |
|
""" |
|
|
|
|
|
if torch.cuda.is_available(): |
|
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16) |
|
pipe = DiffusionPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda") |
|
|
|
|
|
|
|
@spaces.GPU() |
|
def generate_image(prompt, ckpt): |
|
global loaded |
|
print(prompt, ckpt) |
|
|
|
checkpoint = checkpoints[ckpt][0] |
|
num_inference_steps = checkpoints[ckpt][1] |
|
|
|
if loaded != num_inference_steps: |
|
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") |
|
pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, checkpoint), map_location="cuda")) |
|
loaded = num_inference_steps |
|
|
|
if loaded == 1: |
|
results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0, timesteps=[399]) |
|
else: |
|
results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0, timesteps=[999, 749, 499, 249]) |
|
|
|
|
|
return results.images[0] |
|
|
|
|
|
examples = [ |
|
"a cat eating a piece of cheese", |
|
"a ROBOT riding a BLUE horse on Mars, photorealistic", |
|
"Ironman VS Hulk, ultrarealistic", |
|
"a CUTE robot artist painting on an easel", |
|
"Astronaut in a jungle, cold color palette, oil pastel, detailed, 8k", |
|
"An alien holding sign board contain word 'Flash', futuristic, neonpunk", |
|
"Kids going to school, Anime style" |
|
] |
|
|
|
|
|
|
|
|
|
with gr.Blocks(css=CSS) as demo: |
|
gr.HTML("<h1><center>Adobe DMD2🦖</center></h1>") |
|
gr.HTML("<p><center><a href='https://huggingface.co/tianweiy/DMD2'>DMD2</a> text-to-image generation</center></p>") |
|
with gr.Group(): |
|
with gr.Row(): |
|
prompt = gr.Textbox(label='Enter your prompt (English)', scale=8) |
|
ckpt = gr.Dropdown(label='Select inference steps',choices=['1-Step', '4-Step'], value='4-Step', interactive=True) |
|
submit = gr.Button(scale=1, variant='primary') |
|
img = gr.Image(label='DMD2 Generated Image') |
|
gr.Examples( |
|
examples=examples, |
|
inputs=prompt, |
|
outputs=img, |
|
fn=generate_image, |
|
cache_examples=True, |
|
) |
|
|
|
prompt.submit(fn=generate_image, |
|
inputs=[prompt, ckpt], |
|
outputs=img, |
|
) |
|
submit.click(fn=generate_image, |
|
inputs=[prompt, ckpt], |
|
outputs=img, |
|
) |
|
|
|
demo.queue().launch() |