|
|
|
|
|
from __future__ import annotations |
|
|
|
import os |
|
import random |
|
|
|
import gradio as gr |
|
import numpy as np |
|
import torch |
|
from diffusers import AutoencoderKL, DiffusionPipeline |
|
|
|
MAX_SEED = np.iinfo(np.int32).max |
|
MAX_IMG_SIZE = 4096 |
|
|
|
device = torch.device("cpu") |
|
|
|
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32) |
|
pipe = DiffusionPipeline.from_pretrained( |
|
"stabilityai/stable-diffusion-xl-base-1.0", |
|
vae=vae, |
|
torch_dtype=torch.float32, |
|
use_safetensors=True, |
|
variant="fp16", |
|
) |
|
|
|
def random_seed(seed: int, randomize: bool): |
|
if randomize: |
|
seed = random.randint(0, MAX_SEED) |
|
|
|
return seed |
|
|
|
def generate( |
|
prompt: str = "a cat", |
|
seed: int = 0, |
|
width: int = 1024, |
|
height: int = 1024, |
|
): |
|
generator = torch.Generator().manual_seed(seed) |
|
|
|
return pipe( |
|
prompt=prompt, |
|
negative_prompt=None, |
|
prompt_2=None, |
|
negative_prompt_2=None, |
|
width=width, |
|
height=height, |
|
guidance_scale=5.0, |
|
num_inference_steps=10, |
|
generator=generator, |
|
output_type="pil", |
|
).images[0] |
|
|
|
|
|
with gr.Blocks() as instance: |
|
gr.Markdown('# Stable Diffusion') |
|
|
|
with gr.Group(): |
|
prompt = gr.Textbox( |
|
label="Prompt" |
|
) |
|
|
|
seed = gr.Slider( |
|
label="Seed", |
|
minimum=0, |
|
maximum=MAX_SEED, |
|
step=1, |
|
value=0, |
|
) |
|
|
|
is_random_seed = gr.Checkbox(label="Random seed", value=True) |
|
|
|
with gr.Row(): |
|
width = gr.Slider( |
|
label="Width", |
|
minimum=256, |
|
maximum=MAX_IMG_SIZE, |
|
step=32, |
|
value=1024, |
|
) |
|
|
|
height = gr.Slider( |
|
label="Height", |
|
minimum=256, |
|
maximum=MAX_IMG_SIZE, |
|
step=32, |
|
value=1024, |
|
) |
|
|
|
result = gr.Image(label="Result", show_label=False) |
|
|
|
submit = gr.Button("Generate Image") |
|
|
|
gr.on( |
|
triggers=[ |
|
submit.click |
|
|
|
|
|
|
|
], |
|
fn=random_seed, |
|
inputs=[ |
|
seed, |
|
is_random_seed |
|
], |
|
outputs=seed |
|
).then( |
|
fn=generate, |
|
inputs=[ |
|
prompt, |
|
seed, |
|
width, |
|
height, |
|
], |
|
outputs=result |
|
) |
|
|
|
if __name__ == "__main__": |
|
instance.launch() |
|
|
