import space 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", torch_dtype=weight_type) pipe.unet = PeftModel.from_pretrained(pipe.unet, "IDKiro/sdxs-512-dreamshaper-anime") pipe.to(torch_device=device, torch_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 = [ "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k", ] 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)