app.py

"""The Big Sleep — advadnoun's BigGAN x CLIP notebook (Jan 2021), on ZeroGPU.

Engine math is transcribed verbatim from the original Colab (see big_sleep.py):
the layer-wise BigGAN generator, `Pars`, 128 random crops, the latent/class
regularizers and Adam(lr 0.07). This file only adds the UI, live previews,
a timelapse and the notebook's diagnostics panel.
"""
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import time

import spaces  # must precede torch
import gradio as gr
import imageio.v2 as imageio
import numpy as np
import torch

import big_sleep as bs
import clip

HERE = os.path.dirname(os.path.abspath(__file__))

GANS = {"512×512 (the original)": "biggan-deep-512",
        "256×256": "biggan-deep-256",
        "128×128 (fastest)": "biggan-deep-128"}
PERCEPTORS = {"ViT-B/32 (the original)": "ViT-B/32",
              "ViT-B/16 (4× slower)": "ViT-B/16",
              "ViT-L/14 (17× slower)": "ViT-L/14"}
PREVIEW_EVERY = 10
MAX_GPU_SECONDS = 600
# cutouts held in memory at once. B/32 at the notebook's 128 fits in one pass, so the
# default run takes the notebook's exact step; bigger perceptors need accumulation.
CHUNK = {"ViT-B/32": 128, "ViT-B/16": 32, "ViT-L/14": 8}
# cutouts a perceptor can afford before a run stops being watchable
DEFAULT_CUTN = {"ViT-B/32": 128, "ViT-B/16": 48, "ViT-L/14": 16}
# measured ms/iteration on an A100: generator forward+backward, and per cutout through CLIP
GAN_MS = {"biggan-deep-128": 40.0, "biggan-deep-256": 49.0, "biggan-deep-512": 57.0}
CUT_MS = {"ViT-B/32": 1.4, "ViT-B/16": 5.4, "ViT-L/14": 24.5}

print("Loading BigGAN generators...")
MODELS = {v: bs.load_biggan(v).to("cuda") for v in GANS.values()}
print("Loading CLIP perceptors...")
PERS = {v: clip.load(v, device="cpu", jit=False)[0].eval().requires_grad_(False).to("cuda")
        for v in PERCEPTORS.values()}

DIAG_LABELS = {-1: "latent only (class vector zeroed)",
               -2: "class only (fresh random latent)",
               -3: "row 0 broadcast to every layer — plain BigGAN"}


def _step_ms(gan, perceptor, cutn):
    return GAN_MS[GANS[gan]] + CUT_MS[PERCEPTORS[perceptor]] * int(cutn)


def _estimate(iterations, gan, perceptor, cutn):
    return 25 + 1.3 * int(iterations) * _step_ms(gan, perceptor, cutn) / 1000


def _duration(prompt, iterations=500, gan=None, perceptor=None, seed=0, lr=bs.LR,
              cutn=bs.CUTN, timelapse=True):
    gan = gan or list(GANS)[0]
    perceptor = perceptor or list(PERCEPTORS)[0]
    return int(min(MAX_GPU_SECONDS, _estimate(iterations, gan, perceptor, cutn)))


# arg order starts with (prompt, iterations) so gr.Examples can pass just those two
@spaces.GPU(duration=_duration)
def dream(prompt, iterations=500, gan=None, perceptor=None, seed=0, lr=bs.LR,
          cutn=bs.CUTN, timelapse=True):
    gan = gan or list(GANS)[0]
    perceptor = perceptor or list(PERCEPTORS)[0]
    prompt = (prompt or "").strip() or "a cityscape in the style of Van Gogh"
    iterations, cutn, seed = int(iterations), int(cutn), int(seed)
    if seed < 0:
        seed = int(torch.seed() % 2**31)

    model, per = MODELS[GANS[gan]], PERS[PERCEPTORS[perceptor]]
    chunk = CHUNK[PERCEPTORS[perceptor]]
    note = f"{model.side}px · {PERCEPTORS[perceptor]} · {cutn} cutouts · lr {float(lr):g} · seed {seed}"
    if chunk < cutn:
        note += f" · cutouts accumulated {chunk} at a time"

    frames, diagnostics, image = [], [], None
    t0 = rate_t0 = time.time()
    rate_i0 = 0
    for i, img, losses in bs.dream(model, per, prompt, iterations=iterations, seed=seed,
                                  lr=float(lr), cutn=cutn, chunk=chunk,
                                  preview_every=PREVIEW_EVERY):
        if i < 0:
            diagnostics.append((img, DIAG_LABELS[i]))
            continue
        image = img
        if timelapse:
            frames.append(np.asarray(img))
        if i == PREVIEW_EVERY:  # start the rate clock after warmup
            rate_t0, rate_i0 = time.time(), i
        now = time.time()
        rate = (f"{(i - rate_i0) / (now - rate_t0):.1f} it/s"
                if i > rate_i0 else "warming up")
        yield (image, gr.skip(), gr.skip(),
               f"iteration {i}/{iterations} · {rate} · CLIP similarity {-losses[2] / 100:.3f}\n{note}")

    dt = time.time() - rate_t0
    n = max(iterations - rate_i0, 1)
    video = None
    if timelapse and len(frames) > 1:
        video = os.path.join(HERE, "timelapse.mp4")
        imageio.mimsave(video, frames, fps=12, quality=8, macro_block_size=None)
    yield (image, video, diagnostics,
           f"✨ done — {iterations} iterations in {time.time() - t0:.0f}s ({n / dt:.1f} it/s)\n{note}")


CSS = """
.gradio-container { max-width: 1180px !important; margin: 0 auto !important; }
#bs-title { text-align: center; }
"""

with gr.Blocks(title="The Big Sleep") as demo:
    gr.Markdown(
        "# 🛌 The Big Sleep\n"
        "CLIP steers BigGAN toward your words — [advadnoun's original "
        "*BigGANxCLIP* notebook](https://colab.research.google.com/drive/1NCceX2mbiKOSlAd_o7IU7nA9UskKN5WR) "
        "(Ryan Murdock, January 2021), the notebook that started text-to-image as we know it, "
        "running live on ZeroGPU. Nothing is being *sampled*: a batch of latent and class "
        "vectors — **one pair per generator layer** — is optimized until CLIP agrees the picture "
        "matches the prompt.",
        elem_id="bs-title",
    )
    with gr.Row():
        with gr.Column(scale=5):
            prompt = gr.Textbox(label="Prompt", lines=2,
                                value="a cityscape in the style of Van Gogh")
            with gr.Row():
                gan = gr.Dropdown(list(GANS), value=list(GANS)[0], label="BigGAN")
                perceptor = gr.Dropdown(list(PERCEPTORS), value=list(PERCEPTORS)[0],
                                        label="Perceptor")
            iterations = gr.Slider(50, 1500, 500, step=10, label="Iterations")
            with gr.Accordion("Advanced — the notebook's magical numbers", open=False):
                lr = gr.Slider(0.01, 0.2, bs.LR, step=0.005, label="Learning rate")
                cutn = gr.Slider(16, 128, bs.CUTN, step=16, label="Cutouts per iteration")
                seed = gr.Number(0, label="Seed (-1 = random)", precision=0)
                timelapse = gr.Checkbox(True, label="Record timelapse")
            estimate = gr.Markdown()
            with gr.Row():
                run_btn = gr.Button("🛌 Dream", variant="primary", size="lg")
                stop_btn = gr.Button("Stop", variant="stop", size="lg")
        with gr.Column(scale=6):
            out_image = gr.Image(label="Dream", type="pil", format="jpeg", height=480)
            status = gr.Textbox(label="Status", lines=2, interactive=False)
            with gr.Accordion("Timelapse", open=False):
                out_video = gr.Video(label="Timelapse", autoplay=True, loop=True)
            with gr.Accordion("Diagnostics — what each half of the latent is doing", open=False):
                out_diag = gr.Gallery(label="Diagnostics", columns=3, height=240,
                                      object_fit="contain")

    def _show_estimate(iterations, gan, perceptor, cutn):
        secs = _estimate(iterations, gan, perceptor, cutn)
        if secs > MAX_GPU_SECONDS:
            return (f"⚠️ ~{secs / 60:.0f} min of GPU time — over the {MAX_GPU_SECONDS // 60}-minute "
                    "limit, so the run would be cut short. Lower the iterations or the cutouts.")
        return f"≈ {secs:.0f} s of GPU time"

    def _perceptor_changed(perceptor, iterations, gan):
        n = DEFAULT_CUTN[PERCEPTORS[perceptor]]
        return n, _show_estimate(iterations, gan, perceptor, n)

    est_inputs = [iterations, gan, perceptor, cutn]
    for c in (iterations, gan, cutn):
        c.change(_show_estimate, est_inputs, estimate)
    perceptor.change(_perceptor_changed, [perceptor, iterations, gan], [cutn, estimate])
    demo.load(_show_estimate, est_inputs, estimate)

    ev = run_btn.click(dream,
                       [prompt, iterations, gan, perceptor, seed, lr, cutn, timelapse],
                       [out_image, out_video, out_diag, status])
    stop_btn.click(fn=None, cancels=[ev])

    gr.Examples(
        fn=dream,
        examples=[
            ["a cityscape in the style of Van Gogh", 500],
            ["the shadow of a lighthouse falling on a wave", 700],
            ["an ancient library on fire, oil painting", 700],
            ["a demon made of stained glass", 700],
        ],
        inputs=[prompt, iterations],
        outputs=[out_image, out_video, out_diag, status],
        cache_examples=True,
        cache_mode="lazy",
    )

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS)