Spaces:
Paused
Paused
| """ | |
| Stable Diffusion Lab: a teaching Space for BA Design students. | |
| Students can explore how the main knobs of Stable Diffusion 1.5 change an image: | |
| prompt / negative prompt, CFG scale, number of denoising steps, sampler, checkpoint, | |
| seed and resolution. While the image is being generated, a live preview of the | |
| denoising process is streamed to the browser. | |
| How the live preview works | |
| -------------------------- | |
| We run our own (short) denoising loop instead of calling pipe(...). At every step the | |
| U-Net predicts the noise. From that prediction we compute the model's current guess of | |
| the final, clean latent ("x0 prediction") and decode it with TAESD, a tiny and very | |
| fast VAE decoder. The last step is decoded with the full-quality VAE of the checkpoint. | |
| Runs on a dedicated GPU (fp16) but also works on CPU for local testing (slow). | |
| """ | |
| import os | |
| import random | |
| import threading | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from diffusers import ( | |
| AutoencoderKL, | |
| AutoencoderTiny, | |
| DDIMScheduler, | |
| EulerAncestralDiscreteScheduler, | |
| EulerDiscreteScheduler, | |
| StableDiffusionPipeline, | |
| UniPCMultistepScheduler, | |
| ) | |
| from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image | |
| from transformers import CLIPImageProcessor | |
| # ---------------------------------------------------------------------------- | |
| # Configuration (all overridable with Space variables / environment variables) | |
| # ---------------------------------------------------------------------------- | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| DTYPE = getattr(torch, os.getenv("DTYPE", "float16" if DEVICE == "cuda" else "float32")) | |
| SD15_REPO = os.getenv("SD15_REPO", "sd-legacy/stable-diffusion-v1-5") | |
| RV_REPO = os.getenv("RV_REPO", "SG161222/Realistic_Vision_V6.0_B1_noVAE") | |
| RV_FILE = os.getenv("RV_FILE", "Realistic_Vision_V6.0_NV_B1_fp16.safetensors") | |
| # Realistic Vision "noVAE" needs an external VAE (the author recommends sd-vae-ft-mse). | |
| RV_VAE_REPO = os.getenv("RV_VAE_REPO", "stabilityai/sd-vae-ft-mse") | |
| TAESD_REPO = os.getenv("TAESD_REPO", "madebyollin/taesd") | |
| SAFETY_REPO = os.getenv("SAFETY_REPO", "CompVis/stable-diffusion-safety-checker") | |
| USE_SAFETY_CHECKER = os.getenv("SAFETY_CHECKER", "1") == "1" | |
| PRELOAD = os.getenv("PRELOAD", "1") == "1" # load both checkpoints at startup | |
| MAX_CONCURRENT = int(os.getenv("MAX_CONCURRENT", "1")) # generations running at once | |
| MODEL_NAMES = ["Stable Diffusion 1.5", "Realistic Vision V6.0"] | |
| DEFAULT_MODEL = MODEL_NAMES[0] | |
| SAMPLERS = { | |
| "Euler": EulerDiscreteScheduler, | |
| "Euler Ancestral": EulerAncestralDiscreteScheduler, | |
| "DDIM": DDIMScheduler, | |
| "UniPC": UniPCMultistepScheduler, | |
| } | |
| DEFAULTS = dict(steps=20, cfg=7.5, width=512, height=512, seed=42, sampler="Euler") | |
| # ---------------------------------------------------------------------------- | |
| # Model loading | |
| # ---------------------------------------------------------------------------- | |
| _pipes: dict = {} | |
| _taesd = None | |
| _safety = None | |
| _load_lock = threading.Lock() | |
| def _load_sd15() -> StableDiffusionPipeline: | |
| return StableDiffusionPipeline.from_pretrained( | |
| SD15_REPO, | |
| torch_dtype=DTYPE, | |
| variant="fp16", | |
| use_safetensors=True, | |
| safety_checker=None, # we run the safety checker ourselves (see below) | |
| requires_safety_checker=False, | |
| ) | |
| def _load_realistic_vision() -> StableDiffusionPipeline: | |
| vae = AutoencoderKL.from_pretrained(RV_VAE_REPO, torch_dtype=DTYPE) | |
| ckpt = hf_hub_download(RV_REPO, RV_FILE) | |
| return StableDiffusionPipeline.from_single_file( | |
| ckpt, | |
| torch_dtype=DTYPE, | |
| vae=vae, | |
| config=SD15_REPO, # SD 1.5 architecture, so the config files can come from there | |
| safety_checker=None, | |
| requires_safety_checker=False, | |
| ) | |
| _LOADERS = { | |
| "Stable Diffusion 1.5": _load_sd15, | |
| "Realistic Vision V6.0": _load_realistic_vision, | |
| } | |
| def get_pipe(name: str) -> StableDiffusionPipeline: | |
| with _load_lock: | |
| if name not in _pipes: | |
| print(f"Loading {name} ...") | |
| pipe = _LOADERS[name]().to(DEVICE) | |
| pipe.set_progress_bar_config(disable=True) | |
| _pipes[name] = pipe | |
| return _pipes[name] | |
| def get_taesd() -> AutoencoderTiny: | |
| global _taesd | |
| with _load_lock: | |
| if _taesd is None: | |
| _taesd = AutoencoderTiny.from_pretrained(TAESD_REPO, torch_dtype=DTYPE).to(DEVICE) | |
| return _taesd | |
| def get_safety(): | |
| global _safety | |
| with _load_lock: | |
| if _safety is None: | |
| checker = StableDiffusionSafetyChecker.from_pretrained(SAFETY_REPO, torch_dtype=DTYPE).to(DEVICE) | |
| extractor = CLIPImageProcessor.from_pretrained(SAFETY_REPO) | |
| _safety = (checker, extractor) | |
| return _safety | |
| # ---------------------------------------------------------------------------- | |
| # Helpers | |
| # ---------------------------------------------------------------------------- | |
| def to_pil(x: torch.Tensor) -> Image.Image: | |
| """(1, 3, H, W) tensor in [-1, 1] -> PIL image.""" | |
| x = (x.float() / 2 + 0.5).clamp(0, 1)[0].permute(1, 2, 0).cpu().numpy() | |
| return Image.fromarray((x * 255).round().astype(np.uint8)) | |
| def predict_x0(scheduler, x_in: torch.Tensor, eps: torch.Tensor, t) -> torch.Tensor: | |
| """ | |
| The model's current guess of the final clean latent, given the noisy latent and the | |
| predicted noise: x0 = (x_t - sqrt(1 - abar) * eps) / sqrt(abar). | |
| `x_in` is the scheduler-scaled model input, i.e. the sample in "variance preserving" | |
| form. For the Euler samplers abar follows from sigma (abar = 1 / (sigma^2 + 1)), | |
| for DDIM and UniPC it is read from the training noise schedule. | |
| """ | |
| if isinstance(scheduler, (EulerDiscreteScheduler, EulerAncestralDiscreteScheduler)): | |
| sigma = float(scheduler.sigmas[scheduler.step_index]) | |
| abar = 1.0 / (sigma**2 + 1.0) | |
| else: | |
| abar = float(scheduler.alphas_cumprod[int(t)]) | |
| return (x_in.float() - (1.0 - abar) ** 0.5 * eps.float()) / abar**0.5 | |
| def apply_safety_checker(image: Image.Image): | |
| """Returns (image, was_blocked).""" | |
| if not USE_SAFETY_CHECKER: | |
| return image, False | |
| checker, extractor = get_safety() | |
| clip_input = extractor(images=image, return_tensors="pt").pixel_values.to(DEVICE, DTYPE) | |
| arr = np.asarray(image).astype(np.float32)[None] / 255.0 | |
| arr, flagged = checker(images=arr, clip_input=clip_input) | |
| if flagged and flagged[0]: | |
| return Image.fromarray((arr[0] * 255).astype(np.uint8)), True | |
| return image, False | |
| # ---------------------------------------------------------------------------- | |
| # Generation (a generator: every yield updates the UI) | |
| # ---------------------------------------------------------------------------- | |
| def generate(model_name, prompt, negative_prompt, cfg, steps, sampler_name, seed, randomize_seed, width, height): | |
| steps, width, height = int(steps), int(width), int(height) | |
| seed = random.randint(0, 2**32 - 1) if randomize_seed or seed is None else int(seed) | |
| pipe = get_pipe(model_name) | |
| taesd = get_taesd() | |
| # A fresh scheduler per request: schedulers keep internal state, and several | |
| # students may be using the same pipeline. | |
| scheduler = SAMPLERS[sampler_name].from_config(pipe.scheduler.config, timestep_spacing="leading") | |
| scheduler.set_timesteps(steps, device=DEVICE) | |
| # Text conditioning. Batch order is [negative, positive] for classifier-free guidance. | |
| prompt_embeds, negative_embeds = pipe.encode_prompt( | |
| prompt or "", DEVICE, 1, True, negative_prompt=negative_prompt or "" | |
| )[:2] | |
| text_embeds = torch.cat([negative_embeds, prompt_embeds]) | |
| # Seeded starting noise. Created on the CPU so a seed gives the same noise on any machine. | |
| generator = torch.Generator("cpu").manual_seed(seed) | |
| latents = torch.randn((1, pipe.unet.config.in_channels, height // 8, width // 8), generator=generator) | |
| latents = latents.to(DEVICE, DTYPE) * scheduler.init_noise_sigma | |
| step_kwargs = {"generator": generator} if isinstance(scheduler, EulerAncestralDiscreteScheduler) else {} | |
| for i, t in enumerate(scheduler.timesteps): | |
| x_in = scheduler.scale_model_input(torch.cat([latents, latents]), t) | |
| eps = pipe.unet(x_in, t, encoder_hidden_states=text_embeds).sample | |
| eps_uncond, eps_cond = eps.chunk(2) | |
| # Classifier-free guidance: 0 = ignore the prompt, 1 = plain prompt, >1 = push harder. | |
| eps = eps_uncond + cfg * (eps_cond - eps_uncond) | |
| if i < len(scheduler.timesteps) - 1: | |
| x0 = predict_x0(scheduler, x_in.chunk(2)[0], eps, t) | |
| preview = to_pil(taesd.decode(x0.to(DTYPE)).sample) | |
| latents = scheduler.step(eps, t, latents, **step_kwargs).prev_sample | |
| yield preview, seed, f"Step {i + 1} of {steps}: preview of the model's current guess" | |
| else: | |
| latents = scheduler.step(eps, t, latents, **step_kwargs).prev_sample | |
| # Final image with the full VAE of the checkpoint. | |
| vae = pipe.vae | |
| image = to_pil(vae.decode(latents / vae.config.scaling_factor).sample) | |
| image, blocked = apply_safety_checker(image) | |
| if blocked: | |
| gr.Warning("The safety filter blocked this image. Try a different prompt or seed.") | |
| yield image, seed, f"Done: {steps} steps, seed {seed}" + (" (blocked by safety filter)" if blocked else "") | |
| # ---------------------------------------------------------------------------- | |
| # UI | |
| # ---------------------------------------------------------------------------- | |
| EXAMPLES = [ | |
| ["a cozy reading nook with a large window, soft morning light, plants, film photo", "blurry, low quality"], | |
| ["portrait of an old fisherman, weathered face, dramatic side light, 85mm photo", "cartoon, painting, deformed"], | |
| ["a futuristic tram station in the rain at night, neon reflections, wide angle", "text, watermark"], | |
| ["a ceramic teapot shaped like a mountain, studio product photo", ""], | |
| ] | |
| with gr.Blocks(title="Stable Diffusion Lab") as demo: | |
| gr.Markdown( | |
| "# Stable Diffusion Lab\n" | |
| "Explore how text-to-image diffusion works. Change one setting at a time and keep the " | |
| "**seed** fixed to see exactly what that setting does. Watch the image emerge from noise " | |
| "in the live preview." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=5): | |
| model = gr.Dropdown( | |
| MODEL_NAMES, | |
| value=DEFAULT_MODEL, | |
| label="Checkpoint", | |
| info="Stable Diffusion 1.5 is the base model. Realistic Vision V6.0 is a finetune of it " | |
| "that was trained further towards photorealistic images.", | |
| ) | |
| prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe what you want to see") | |
| negative = gr.Textbox( | |
| label="Negative prompt", | |
| lines=2, | |
| placeholder="What the image should avoid, e.g. blurry, low quality, extra fingers", | |
| ) | |
| with gr.Row(): | |
| sampler = gr.Dropdown( | |
| list(SAMPLERS), | |
| value=DEFAULTS["sampler"], | |
| label="Sampler", | |
| info="The algorithm that walks from noise to image. Euler Ancestral adds fresh noise " | |
| "at every step, so it never fully settles.", | |
| ) | |
| steps = gr.Slider( | |
| 4, 48, value=DEFAULTS["steps"], step=1, label="Denoising steps", | |
| info="More steps = finer refinement, but slower.", | |
| ) | |
| cfg = gr.Slider( | |
| 0.0, 15.0, value=DEFAULTS["cfg"], step=0.1, label="CFG scale", | |
| info="How strongly the prompt is enforced. 0 ignores the prompt, 1 follows it plainly, " | |
| "high values get rigid and oversaturated.", | |
| ) | |
| with gr.Row(): | |
| width = gr.Slider( | |
| 256, 768, value=DEFAULTS["width"], step=64, label="Width (px)", | |
| info="Both models were trained at 512 px. Much larger sizes tend to repeat objects.", | |
| ) | |
| height = gr.Slider(256, 768, value=DEFAULTS["height"], step=64, label="Height (px)") | |
| with gr.Row(): | |
| seed = gr.Number( | |
| value=DEFAULTS["seed"], precision=0, minimum=0, maximum=2**32 - 1, label="Seed", | |
| info="Same seed + same settings = same image.", | |
| ) | |
| randomize = gr.Checkbox(value=False, label="Random seed", info="Pick a new seed on every run.") | |
| run = gr.Button("Generate", variant="primary") | |
| with gr.Column(scale=5): | |
| output = gr.Image(label="Result (live preview while generating)", type="pil", interactive=False) | |
| status = gr.Markdown() | |
| gr.Examples(EXAMPLES, inputs=[prompt, negative], label="Example prompts") | |
| inputs = [model, prompt, negative, cfg, steps, sampler, seed, randomize, width, height] | |
| outputs = [output, seed, status] | |
| for trigger in (run.click, prompt.submit): | |
| trigger( | |
| generate, inputs, outputs, | |
| show_progress="hidden", concurrency_limit=MAX_CONCURRENT, api_name="generate" if trigger == run.click else False, | |
| ) | |
| demo.queue(max_size=40) | |
| if __name__ == "__main__": | |
| if PRELOAD: | |
| for _name in MODEL_NAMES: | |
| get_pipe(_name) | |
| get_taesd() | |
| if USE_SAFETY_CHECKER: | |
| get_safety() | |
| demo.launch() | |