File size: 5,306 Bytes
d9c331d 03055f8 645f954 783125d f1c0b37 03055f8 d9c331d f1c0b37 b4fabad f3eb48f 783125d b4fabad c8b1bf4 74c2772 783125d 74c2772 694f32d 783125d 74c2772 33e0372 783125d 694f32d 783125d 74c2772 783125d 694f32d 03055f8 bc9a8f6 783125d bc9a8f6 56a8823 29593cd bc9a8f6 03055f8 bc9a8f6 03055f8 783125d 03055f8 74c2772 03055f8 74c2772 03055f8 74c2772 03055f8 74c2772 03055f8 74c2772 03055f8 74c2772 03055f8 74c2772 03055f8 f1c0b37 d9c331d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
from fastapi import FastAPI
import gradio as gr
import torch
from sd import pipeline
from sd import model_loader
from transformers import AutoTokenizer
from diffusers import StableDiffusionPipeline
app = FastAPI()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
weights_url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt"
models = model_loader.from_pretrained(weights_url, device)
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_safetensors=True)
pipe = pipe.to(device)
MIN_IMAGE_SIZE = 256
MAX_IMAGE_SIZE = 1024
MAX_SEED = 2147483647 # 2^31 - 1
def generate_image(prompt, negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, model, width, height):
if randomize_seed:
seed = torch.randint(0, MAX_SEED, (1,)).item()
generator = torch.Generator(device=device).manual_seed(seed)
if model == "from-scratch":
image = pipeline.generate(
prompt=prompt,
uncond_prompt=negative_prompt,
input_image=None,
strength=0.9,
cfg_scale=guidance_scale,
n_inference_steps=num_inference_steps,
width=width,
height=height,
generator=generator,
device=device,
idle_device="cpu",
models=models,
tokenizer=tokenizer,
)
else:
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=width,
height=height,
generator=generator,
).images[0]
return image
css="""
#col-container {
margin: 0 auto;
max-width: 520px;
}
"""
md = """
# Text-to-Image: Stable Diffusion from Scratch
### By [Nazareno Amidolare](https://kepler296e.github.io/)
Using **Docker**, **FastAPI**, **PyTorch** and **Gradio**.
### References
- [Coding Stable Diffusion from scratch in PyTorch](https://www.youtube.com/watch?v=ZBKpAp_6TGI&ab_channel=UmarJamil)
- [Hugging Space Diffusers](https://github.com/huggingface/diffusers/)
Currently running on a **CPU** (≈20 minutes per image).
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(md)
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)
result = gr.Image(label="Result", show_label=False)
with gr.Accordion("Advanced Settings", open=False):
with gr.Row():
model = gr.Dropdown(
label="Model",
choices=["from-scratch", "runwayml/stable-diffusion-v1-5"],
value="from-scratch",
interactive=True,
)
with gr.Row():
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=42,
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
negative_prompt = gr.Text(
label="Negative prompt",
max_lines=1,
placeholder="Enter a negative prompt",
visible=False,
value="",
)
with gr.Row():
width = gr.Slider(
label="Width",
minimum=MIN_IMAGE_SIZE,
maximum=MAX_IMAGE_SIZE,
step=32,
value=512,
)
height = gr.Slider(
label="Height",
minimum=MIN_IMAGE_SIZE,
maximum=MAX_IMAGE_SIZE,
step=32,
value=512,
)
with gr.Row():
guidance_scale = gr.Slider(
label="Guidance scale",
minimum=1.0,
maximum=14.0,
step=0.1,
value=8.0,
)
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=50,
step=1,
value=50,
)
run_button.click(
fn = generate_image,
inputs = [prompt, negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, model, width, height],
outputs = [result]
)
app = gr.mount_gradio_app(app, demo, "/")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000) |