| import spaces |
| import base64 |
| from io import BytesIO |
|
|
| import gradio as gr |
| import PIL.Image |
| import torch |
|
|
| from diffusers import StableDiffusionPipeline, AutoencoderKL, AutoencoderTiny |
| from peft import PeftModel |
|
|
| device = "cuda" |
| weight_type = torch.float16 |
|
|
| pipe = StableDiffusionPipeline.from_pretrained("IDKiro/sdxs-512-dreamshaper") |
| pipe.unet = PeftModel.from_pretrained(pipe.unet, "IDKiro/sdxs-512-dreamshaper-anime") |
| pipe.unet.merge_and_unload() |
| pipe.to(device, dtype=weight_type) |
|
|
| vae_tiny = AutoencoderTiny.from_pretrained( |
| "IDKiro/sdxs-512-dreamshaper", subfolder="vae" |
| ) |
| vae_tiny.to(device, dtype=weight_type) |
|
|
| vae_large = AutoencoderKL.from_pretrained( |
| "IDKiro/sdxs-512-dreamshaper", subfolder="vae_large" |
| ) |
| vae_tiny.to(device, dtype=weight_type) |
|
|
|
|
| def pil_image_to_data_url(img, format="PNG"): |
| buffered = BytesIO() |
| img.save(buffered, format=format) |
| img_str = base64.b64encode(buffered.getvalue()).decode() |
| return f"data:image/{format.lower()};base64,{img_str}" |
|
|
|
|
| @spaces.GPU |
| def run( |
| prompt: str, |
| device_type="GPU", |
| vae_type=None, |
| param_dtype="torch.float16", |
| ) -> PIL.Image.Image: |
| if vae_type == "tiny vae": |
| pipe.vae = vae_tiny |
| elif vae_type == "large vae": |
| pipe.vae = vae_large |
|
|
| if device_type == "CPU": |
| device = "cpu" |
| param_dtype = "torch.float32" |
| else: |
| device = "cuda" |
|
|
| pipe.to( |
| torch_device=device, |
| torch_dtype=torch.float16 if param_dtype == "torch.float16" else torch.float32, |
| ) |
|
|
| result = pipe( |
| prompt=prompt, |
| guidance_scale=0.0, |
| num_inference_steps=1, |
| output_type="pil", |
| ).images[0] |
|
|
| result_url = pil_image_to_data_url(result) |
|
|
| return (result, result_url) |
|
|
|
|
| examples = [ |
| "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece", |
| ] |
|
|
| with gr.Blocks(css="style.css") as demo: |
| gr.Markdown("# SDXS-512-DreamShaper-Anime") |
| gr.Markdown("[SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions](https://arxiv.org/abs/2403.16627) | [GitHub](https://github.com/IDKiro/sdxs)") |
| with gr.Group(): |
| with gr.Row(): |
| with gr.Column(min_width=685): |
| with gr.Row(): |
| prompt = gr.Text( |
| label="Prompt", |
| show_label=False, |
| max_lines=1, |
| placeholder="Enter your prompt", |
| container=False, |
| ) |
| run_button = gr.Button("Run", scale=0) |
|
|
| device_choices = ["GPU", "CPU"] |
| device_type = gr.Radio( |
| device_choices, |
| label="Device", |
| value=device_choices[0], |
| interactive=True, |
| info="Thanks to the community for the GPU!", |
| ) |
|
|
| vae_choices = ["tiny vae", "large vae"] |
| vae_type = gr.Radio( |
| vae_choices, |
| label="Image Decoder Type", |
| value=vae_choices[0], |
| interactive=True, |
| info="To save GPU memory, use tiny vae. For better quality, use large vae.", |
| ) |
|
|
| dtype_choices = ["torch.float16", "torch.float32"] |
| param_dtype = gr.Radio( |
| dtype_choices, |
| label="torch.weight_type", |
| value=dtype_choices[0], |
| interactive=True, |
| info="To save GPU memory, use torch.float16. For better quality, use torch.float32.", |
| ) |
|
|
| download_output = gr.Button( |
| "Download output", elem_id="download_output" |
| ) |
|
|
| with gr.Column(min_width=512): |
| result = gr.Image( |
| label="Result", |
| height=512, |
| width=512, |
| elem_id="output_image", |
| show_label=False, |
| show_download_button=True, |
| ) |
|
|
| gr.Examples(examples=examples, inputs=prompt, outputs=result, fn=run) |
|
|
| demo.load(None, None, None) |
|
|
| inputs = [prompt, device_type, vae_type, param_dtype] |
| outputs = [result, download_output] |
| prompt.submit(fn=run, inputs=inputs, outputs=outputs) |
| run_button.click(fn=run, inputs=inputs, outputs=outputs) |
|
|
| if __name__ == "__main__": |
| demo.queue().launch(debug=True) |
|
|
