IDKiro's picture
Upload 3 files
1600949 verified
raw history blame
No virus
4.31 kB
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 = "cpu" # Linux & Windows
weight_type = torch.float32 # torch.float16 works as well, but pictures seem to be a bit worse
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}"
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 (only CPU now)")
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='Only CPU now.')
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)