Spaces:

hugging-apps
/

rads-diffusion-steering

Running on Zero

Upload app.py with huggingface_hub

b840066 verified about 16 hours ago

20.9 kB

import spaces  # ZeroGPU: must be imported before anything that touches CUDA

import random
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal

import gradio as gr
from diffusers import DDIMScheduler, StableDiffusionPipeline
from huggingface_hub import hf_hub_download

# ===========================================================================
# Vendored model definitions — copied 1:1 from the official RADS repository
# (https://github.com/S-Karnik/rads-paper-repo, MIT license):
#   - train/sac_steer_lib.py      -> SharedEncoder, Actor (inference only)
#   - clip_autoencoder/model.py   -> CLIPVAE
# The training-only critics (Critic / SafetyCritic) are not used at sampling
# time: with evaluate=True, select_action only runs encoder + actor.
# ===========================================================================


def weights_init_(m):
    if isinstance(m, nn.Linear):
        torch.nn.init.xavier_uniform_(m.weight, gain=1)
        torch.nn.init.constant_(m.bias, 0)


class SharedEncoder(nn.Module):
    """Encodes (denoising latent, prompt latent, normalized timestep) -> features."""

    def __init__(self, caption_dim, time_embed_dim=16, encoded_dim=256, latent_resolution=64):
        super(SharedEncoder, self).__init__()
        self.conv1 = nn.Conv2d(4, 32, kernel_size=3, padding=1)
        self.bn1 = nn.GroupNorm(8, 32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.GroupNorm(8, 64)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.GroupNorm(8, 128)
        final_dim = latent_resolution // 8
        self.flat_img_size = 128 * final_dim * final_dim
        self.fc_img = nn.Linear(self.flat_img_size, 256)
        self.flat_cap_size = 2 * caption_dim
        self.fc_cap = nn.Linear(self.flat_cap_size, 256)
        self.time_embed_dim = time_embed_dim
        self.time_mlp = nn.Sequential(
            nn.Linear(1, time_embed_dim),
            nn.ReLU(),
            nn.Linear(time_embed_dim, time_embed_dim),
            nn.ReLU(),
        )
        self.fc_joint = nn.Linear(512 + time_embed_dim, encoded_dim)
        self.ln_joint = nn.LayerNorm(encoded_dim)
        self.apply(weights_init_)

    def forward(self, img, cap, timestep):
        x = F.relu(self.bn1(self.conv1(img)))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.bn3(self.conv3(x)))
        x = F.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)
        x_img = F.relu(self.fc_img(x))
        if cap.dim() == 3:
            x_c = cap.view(cap.size(0), -1)
            x_cap = F.relu(self.fc_cap(x_c))
        else:
            raise ValueError(f"Expected 3D caption tensor [B, 2, 64], got {cap.shape}")
        if not torch.is_tensor(timestep):
            timestep = torch.tensor(timestep, device=x_img.device, dtype=torch.float32)
        timestep = timestep.to(dtype=torch.float32)

        if timestep.dim() == 0:
            timestep = timestep.reshape(1, 1).expand(x_img.size(0), 1)
        elif timestep.dim() == 1:
            if timestep.size(0) == 1:
                timestep = timestep.expand(x_img.size(0)).unsqueeze(1)
            else:
                timestep = timestep.unsqueeze(1)

        t_emb = self.time_mlp(timestep)
        x_joint = torch.cat([x_img, x_cap, t_emb], dim=1)
        return F.relu(self.ln_joint(self.fc_joint(x_joint)))


class Actor(nn.Module):
    """SAC actor; sample() -> (action, log_prob, deterministic tanh-mean)."""

    def __init__(self, hidden_dim, action_dim=None, action_space=None):
        super(Actor, self).__init__()
        self.linear1 = nn.Linear(hidden_dim, 256)
        self.ln1 = nn.LayerNorm(256)
        self.linear2 = nn.Linear(256, 256)
        self.mean_linear = nn.Linear(256, action_dim)
        self.log_std_linear = nn.Linear(256, action_dim)
        self.apply(weights_init_)
        if action_space is None:
            self.action_scale = torch.tensor(1.)
            self.action_bias = torch.tensor(0.)
        else:
            self.action_scale = torch.FloatTensor((action_space.high - action_space.low) / 2.)
            self.action_bias = torch.FloatTensor((action_space.high + action_space.low) / 2.)

    def forward(self, state_features):
        x = F.relu(self.ln1(self.linear1(state_features)))
        x = F.relu(self.linear2(x))
        mean = self.mean_linear(x)
        log_std = self.log_std_linear(x)
        log_std = torch.clamp(log_std, min=-6, max=2)
        return mean, log_std

    def sample(self, state_features):
        mean, log_std = self.forward(state_features)
        std = log_std.exp()
        normal = Normal(mean, std)
        x_t = normal.rsample()
        y_t = torch.tanh(x_t)
        action = y_t * self.action_scale + self.action_bias
        log_prob = normal.log_prob(x_t)
        log_prob -= torch.log(self.action_scale * (1 - y_t.pow(2)) + 1e-6)
        log_prob = log_prob.sum(1, keepdim=True)
        mean = torch.tanh(mean) * self.action_scale + self.action_bias
        return action, log_prob, mean


class CLIPVAE(nn.Module):
    """Bottleneck autoencoder over CLIP caption embeddings (77 x 768 <-> 64)."""

    def __init__(self, latent_dim=64, seq_len=77, embed_dim=768, n_layers=4, n_heads=8, dropout=0.1, decoder_final_norm=True):
        super().__init__()
        self.seq_len = seq_len
        self.embed_dim = embed_dim

        self.pos_embed = nn.Parameter(torch.randn(1, seq_len, embed_dim) * 0.02)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=n_heads, dim_feedforward=2048,
            dropout=dropout, batch_first=True, norm_first=True
        )
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)

        self.bottleneck_norm = nn.LayerNorm(seq_len * embed_dim)

        self.fc_mu = nn.Linear(seq_len * embed_dim, latent_dim)
        self.fc_var = nn.Linear(seq_len * embed_dim, latent_dim)
        self.fc_decoder_input = nn.Linear(latent_dim, seq_len * embed_dim)

        decoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=n_heads, dim_feedforward=2048,
            dropout=dropout, batch_first=True, norm_first=True
        )
        self.transformer_decoder = nn.TransformerEncoder(decoder_layer, num_layers=n_layers)

        self.decoder_final_norm = nn.LayerNorm(embed_dim) if decoder_final_norm else None

    def reparameterize(self, mu, logvar):
        logvar = torch.clamp(logvar, min=-6, max=6)
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def encode(self, x):
        batch_size = x.shape[0]
        x = x + self.pos_embed
        x_encoded = self.transformer_encoder(x)

        x_flat = x_encoded.reshape(batch_size, -1)
        x_flat = self.bottleneck_norm(x_flat)

        mu = self.fc_mu(x_flat)
        logvar = self.fc_var(x_flat)
        return mu, logvar

    def decode(self, z):
        batch_size = z.shape[0]
        x_decode = self.fc_decoder_input(z).view(batch_size, self.seq_len, self.embed_dim)
        x_decode = x_decode + self.pos_embed

        out = self.transformer_decoder(x_decode)

        out = self.decoder_final_norm(out) if self.decoder_final_norm is not None else out
        return out

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        out = self.decode(z)
        return out, mu, logvar


# ===========================================================================
# Constants (paper configuration: SD v1.4 + DDIM, T = 50, g = 7.5)
# ===========================================================================

CKPT_REPO = "S-Karnik/rads-model-checkpoints"
VAE_FILE = "clip_vae/step=45000.pth"  # CLIP caption-embedding autoencoder (SD v1.4 / RealVis)
TRAIN_SEEDS = [41, 42, 43, 44, 45]    # the 5 SAC training seeds reported in the paper
DEFAULT_TRAIN_SEED = 42
CAPTION_DIM = 64
TIME_EMBED_DIM = 16
BASE_MODEL = "CompVis/stable-diffusion-v1-4"

device = "cuda"


def _load_state_dict(path: str):
    """Load a torch state dict, preferring the safe weights-only path."""
    try:
        return torch.load(path, map_location="cpu", weights_only=True)
    except Exception:
        return torch.load(path, map_location="cpu", weights_only=False)


def _load_rads(train_seed: int):
    """Load the SAC encoder + actor checkpoints trained for `train_seed`."""
    prefix = f"RADS_SDV1_DDIM/checkpoints_SDV1_DDIM_seed_{train_seed}"
    encoder = SharedEncoder(CAPTION_DIM, time_embed_dim=TIME_EMBED_DIM, latent_resolution=64)
    encoder.load_state_dict(_load_state_dict(hf_hub_download(CKPT_REPO, f"{prefix}/sac_encoder_best.pth")))
    encoder.eval()

    actor = Actor(256, action_dim=CAPTION_DIM, action_space=None)
    actor.load_state_dict(_load_state_dict(hf_hub_download(CKPT_REPO, f"{prefix}/sac_actor_best.pth")))
    actor.eval()
    return encoder, actor


# ===========================================================================
# Model loading (module scope, eager, on the GPU — ZeroGPU defers the move)
# ===========================================================================

print("[RADS] loading Stable Diffusion v1.4 (fp16, DDIM)...")
scheduler = DDIMScheduler.from_pretrained(BASE_MODEL, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(
    BASE_MODEL,
    scheduler=scheduler,
    variant="fp16",
    torch_dtype=torch.float16,
    safety_checker=None,
    requires_safety_checker=False,
)
pipe.set_progress_bar_config(disable=True)
pipe = pipe.to(device)

print("[RADS] loading RADS steering components...")
clip_vae = CLIPVAE(latent_dim=CAPTION_DIM, decoder_final_norm=False)
clip_vae.load_state_dict(_load_state_dict(hf_hub_download(CKPT_REPO, VAE_FILE)))
clip_vae.eval()
clip_vae = clip_vae.to(device)

agents = {}
for train_seed in TRAIN_SEEDS:
    enc, act = _load_rads(train_seed)
    agents[train_seed] = (enc.to(device), act.to(device))
print(f"[RADS] ready — loaded SAC agents for train seeds {TRAIN_SEEDS}")


# ===========================================================================
# Inference — replicates utils/run/generate.py (test_step, "sac" branch) from
# the official repo 1:1: per-step SAC action in the 64-d caption latent space,
# steered embedding replaces the conditional branch of classifier-free guidance.
# ===========================================================================


def _select_action(encoder, actor, latents, current_prompt_latents, t_norm):
    """Deterministic (evaluate=True) SAC action for the current state."""
    with torch.no_grad():
        state_img = latents.to(device=device, dtype=torch.float32)
        state_cap = current_prompt_latents.to(device=device, dtype=torch.float32)
        if state_img.dim() == 3:
            state_img = state_img.unsqueeze(0)
        if state_cap.dim() == 2:
            state_cap = state_cap.unsqueeze(0)
        features = encoder(state_img, state_cap, t_norm)
        _, _, action = actor.sample(features)  # evaluate=True -> deterministic mean
    return action.detach()


def _generate_one(prompt, seed, num_steps, guidance, steer_scale, train_seed,
                  progress=None, prog_range=(0.0, 1.0)):
    """Run one SD v1.4 DDIM denoising trajectory; steer it with RADS if scale > 0.
    Returns a single PIL image.
    """
    encoder, actor = agents[int(train_seed)]

    # --- prepare (mirrors wah's prepare_embeds / prepare_timesteps / prepare_latents)
    prompt_embeds, negative_embeds = pipe.encode_prompt(
        prompt=prompt,
        device=device,
        num_images_per_prompt=1,
        do_classifier_free_guidance=True,
        negative_prompt="",
    )
    embeds = torch.cat([negative_embeds, prompt_embeds], dim=0)  # [2, 77, 768] fp16 (neg first)

    generator = torch.Generator(device=device).manual_seed(int(seed))
    pipe.scheduler.set_timesteps(num_steps, device=device)
    timesteps = pipe.scheduler.timesteps
    latents = pipe.prepare_latents(
        1, pipe.unet.config.in_channels, 512, 512,
        embeds.dtype, device, generator, None,
    )

    steer = steer_scale > 0.0
    if steer:
        embeds32 = embeds.to(dtype=torch.float32)
        with torch.no_grad():
            # caption latents fed to the policy: [B, 2, 64] = (neg, pos)
            prompt_mu, _ = clip_vae.encode(embeds32)
            neg_mu, pos_mu = prompt_mu.chunk(2)
            current_prompt_latents = torch.stack([neg_mu, pos_mu], dim=1)
            # mu of the conditional embedding (constant across steps; the reference
            # recomputes it every step from the same input, which is identical)
            cond_mu, _ = clip_vae.encode(embeds32.chunk(2)[1])
        neg_embeds = embeds.chunk(2)[0]

    with torch.no_grad():
        for i, t in enumerate(timesteps):
            if steer:
                action_latent = _select_action(
                    encoder, actor, latents, current_prompt_latents, i / num_steps
                )
                z_steered = cond_mu + (action_latent * steer_scale)
                steered_cond_embedding = clip_vae.decode(z_steered)
                if steered_cond_embedding.dim() == 4:
                    steered_cond_embedding = steered_cond_embedding.squeeze(1)
                steered_embeds = torch.cat(
                    [neg_embeds, steered_cond_embedding.to(dtype=latents.dtype)], dim=0
                )
                step_embeds = steered_embeds
            else:
                step_embeds = embeds

            # --- predict_noise (classifier-free guidance: one batched forward pass)
            latent_model_input = torch.cat([latents] * 2)
            noise_pred = pipe.unet(
                latent_model_input, t, encoder_hidden_states=step_embeds, return_dict=False
            )[0]

            # --- perform_guidance
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance * (noise_pred_text - noise_pred_uncond)

            # --- step (DDIM, eta = 0)
            latents = pipe.scheduler.step(
                noise_pred, t, latents, generator=generator, return_dict=False
            )[0]

            if progress is not None:
                lo, hi = prog_range
                progress(lo + (i + 1) / num_steps * (hi - lo))

    # --- decode
    with torch.no_grad():
        decoded = pipe.vae.decode(
            latents / pipe.vae.config.scaling_factor, return_dict=False, generator=generator
        )[0]
        images = pipe.image_processor.postprocess(
            decoded, output_type="pil", do_denormalize=[True] * decoded.shape[0]
        )
    return images[0]


@spaces.GPU(duration=20)
def generate(prompt, steer_scale=1.0, train_seed=DEFAULT_TRAIN_SEED, seed=0,
             randomize_seed=False, num_steps=50, guidance_scale=7.5,
             progress=gr.Progress(track_tqdm=True)):
    """Generate a baseline and a RADS-steered image for the same prompt and seed.
    Args:
        prompt: text prompt for Stable Diffusion v1.4.
        steer_scale: RADS steering strength (LATENT_SCALE in the paper's code);
            0 disables steering and reproduces the plain SD v1.4 baseline.
        train_seed: which of the 5 published SAC training seeds to use.
        seed: initial latent seed (shared by both images, so the pair differs
            only through steering).
        randomize_seed: draw a fresh random seed for this run.
        num_steps: number of DDIM denoising steps (paper default: 50).
        guidance_scale: classifier-free guidance scale (paper default: 7.5).
    Returns:
        (baseline image, RADS-steered image, seed used)
    """
    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)
    train_seed = int(train_seed)

    t0 = time.perf_counter()
    baseline = _generate_one(
        prompt, seed, int(num_steps), float(guidance_scale), 0.0, train_seed,
        progress=progress, prog_range=(0.0, 0.5),
    )
    t1 = time.perf_counter()
    steered = _generate_one(
        prompt, seed, int(num_steps), float(guidance_scale), float(steer_scale), train_seed,
        progress=progress, prog_range=(0.5, 1.0),
    )
    t2 = time.perf_counter()
    print(f"[RADS] prompt={prompt!r} seed={seed} train_seed={train_seed} "
          f"scale={steer_scale} baseline={t1 - t0:.2f}s steered={t2 - t1:.2f}s")
    return baseline, steered, seed


# ===========================================================================
# UI
# ===========================================================================

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["satisfied happy scientist"],
    ["fire us army garden flag single"],
    ["The No Limits Business Woman Podcast"],
    ["train that glorious with st robin western rug"],
    ["gold colorful transit industry protected sign mosaic rug"],
    ["Michael Fassbender to Star In Assassin's Creed Movie"],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # RADS: Steering Away from Memorization in Text-to-Image Diffusion
            **RADS** (Reachability-constrained reinforcement learning for Diffusion Steering) is an
            inference-time controller that steers **Stable Diffusion v1.4** away from reproducing
            memorized training images. A SAC policy, trained offline against a reachability-based
            safety critic, predicts a small perturbation in a learned 64-d caption-embedding latent
            space at every denoising step — the diffusion backbone itself is never modified.
            This demo generates the **same prompt and seed** twice: plain SD v1.4 (left) vs.
            RADS-steered (right). Both images share the identical initial latent, so any
            difference comes purely from steering. Prompts below are memorization-prone prompts
            from the paper's evaluation sets (Webster 2023 / MemBench).
            [[Paper](https://huggingface.co/papers/2603.00140)] ·
            [Project page](https://s-karnik.github.io/rads-memorization-project-page) ·
            [Code](https://github.com/S-Karnik/rads-paper-repo) ·
            [Checkpoints](https://huggingface.co/S-Karnik/rads-model-checkpoints)
            """
        )

        prompt = gr.Textbox(
            label="Prompt",
            placeholder="A memorization-prone prompt, e.g. 'satisfied happy scientist'",
            lines=2,
        )
        with gr.Row():
            steer_scale = gr.Slider(
                0.0, 2.0, value=1.0, step=0.05,
                label="Steering strength (0 = RADS off, 1 = paper default)",
                info="Scales the SAC action added to the caption-embedding latent (LATENT_SCALE).",
            )
            run_button = gr.Button("Generate", variant="primary", scale=1)

        with gr.Row():
            baseline_image = gr.Image(
                label="Baseline — SD v1.4 (no steering)", type="pil", height=420
            )
            steered_image = gr.Image(
                label="RADS-steered — same seed", type="pil", height=420
            )

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                train_seed = gr.Dropdown(
                    choices=[str(s) for s in TRAIN_SEEDS],
                    value=str(DEFAULT_TRAIN_SEED),
                    label="SAC training seed",
                    info="The paper reports 5 independently trained SAC agents (seeds 41–45).",
                )
                num_steps = gr.Slider(
                    10, 100, value=50, step=1, label="DDIM steps (paper: 50)"
                )
                guidance_scale = gr.Slider(
                    1.0, 15.0, value=7.5, step=0.5, label="Guidance scale (paper: 7.5)"
                )
            with gr.Row():
                seed = gr.Slider(
                    0, 2**31 - 1, value=0, step=1, label="Seed"
                )
                randomize_seed = gr.Checkbox(
                    value=False, label="Randomize seed (both images share it)"
                )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[prompt],
            outputs=[baseline_image, steered_image, seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    inputs = [prompt, steer_scale, train_seed, seed, randomize_seed, num_steps, guidance_scale]
    outputs = [baseline_image, steered_image, seed]
    run_button.click(fn=generate, inputs=inputs, outputs=outputs)

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

	import spaces # ZeroGPU: must be imported before anything that touches CUDA

	import random
	import time

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.distributions import Normal

	import gradio as gr
	from diffusers import DDIMScheduler, StableDiffusionPipeline
	from huggingface_hub import hf_hub_download

	# ===========================================================================
	# Vendored model definitions — copied 1:1 from the official RADS repository
	# (https://github.com/S-Karnik/rads-paper-repo, MIT license):
	# - train/sac_steer_lib.py -> SharedEncoder, Actor (inference only)
	# - clip_autoencoder/model.py -> CLIPVAE
	# The training-only critics (Critic / SafetyCritic) are not used at sampling
	# time: with evaluate=True, select_action only runs encoder + actor.
	# ===========================================================================


	def weights_init_(m):
	if isinstance(m, nn.Linear):
	torch.nn.init.xavier_uniform_(m.weight, gain=1)
	torch.nn.init.constant_(m.bias, 0)


	class SharedEncoder(nn.Module):
	"""Encodes (denoising latent, prompt latent, normalized timestep) -> features."""

	def __init__(self, caption_dim, time_embed_dim=16, encoded_dim=256, latent_resolution=64):
	super(SharedEncoder, self).__init__()
	self.conv1 = nn.Conv2d(4, 32, kernel_size=3, padding=1)
	self.bn1 = nn.GroupNorm(8, 32)
	self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
	self.bn2 = nn.GroupNorm(8, 64)
	self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
	self.bn3 = nn.GroupNorm(8, 128)
	final_dim = latent_resolution // 8
	self.flat_img_size = 128 * final_dim * final_dim
	self.fc_img = nn.Linear(self.flat_img_size, 256)
	self.flat_cap_size = 2 * caption_dim
	self.fc_cap = nn.Linear(self.flat_cap_size, 256)
	self.time_embed_dim = time_embed_dim
	self.time_mlp = nn.Sequential(
	nn.Linear(1, time_embed_dim),
	nn.ReLU(),
	nn.Linear(time_embed_dim, time_embed_dim),
	nn.ReLU(),
	)
	self.fc_joint = nn.Linear(512 + time_embed_dim, encoded_dim)
	self.ln_joint = nn.LayerNorm(encoded_dim)
	self.apply(weights_init_)

	def forward(self, img, cap, timestep):
	x = F.relu(self.bn1(self.conv1(img)))
	x = F.max_pool2d(x, 2)
	x = F.relu(self.bn2(self.conv2(x)))
	x = F.max_pool2d(x, 2)
	x = F.relu(self.bn3(self.conv3(x)))
	x = F.max_pool2d(x, 2)
	x = x.view(x.size(0), -1)
	x_img = F.relu(self.fc_img(x))
	if cap.dim() == 3:
	x_c = cap.view(cap.size(0), -1)
	x_cap = F.relu(self.fc_cap(x_c))
	else:
	raise ValueError(f"Expected 3D caption tensor [B, 2, 64], got {cap.shape}")
	if not torch.is_tensor(timestep):
	timestep = torch.tensor(timestep, device=x_img.device, dtype=torch.float32)
	timestep = timestep.to(dtype=torch.float32)

	if timestep.dim() == 0:
	timestep = timestep.reshape(1, 1).expand(x_img.size(0), 1)
	elif timestep.dim() == 1:
	if timestep.size(0) == 1:
	timestep = timestep.expand(x_img.size(0)).unsqueeze(1)
	else:
	timestep = timestep.unsqueeze(1)

	t_emb = self.time_mlp(timestep)
	x_joint = torch.cat([x_img, x_cap, t_emb], dim=1)
	return F.relu(self.ln_joint(self.fc_joint(x_joint)))


	class Actor(nn.Module):
	"""SAC actor; sample() -> (action, log_prob, deterministic tanh-mean)."""

	def __init__(self, hidden_dim, action_dim=None, action_space=None):
	super(Actor, self).__init__()
	self.linear1 = nn.Linear(hidden_dim, 256)
	self.ln1 = nn.LayerNorm(256)
	self.linear2 = nn.Linear(256, 256)
	self.mean_linear = nn.Linear(256, action_dim)
	self.log_std_linear = nn.Linear(256, action_dim)
	self.apply(weights_init_)
	if action_space is None:
	self.action_scale = torch.tensor(1.)
	self.action_bias = torch.tensor(0.)
	else:
	self.action_scale = torch.FloatTensor((action_space.high - action_space.low) / 2.)
	self.action_bias = torch.FloatTensor((action_space.high + action_space.low) / 2.)

	def forward(self, state_features):
	x = F.relu(self.ln1(self.linear1(state_features)))
	x = F.relu(self.linear2(x))
	mean = self.mean_linear(x)
	log_std = self.log_std_linear(x)
	log_std = torch.clamp(log_std, min=-6, max=2)
	return mean, log_std

	def sample(self, state_features):
	mean, log_std = self.forward(state_features)
	std = log_std.exp()
	normal = Normal(mean, std)
	x_t = normal.rsample()
	y_t = torch.tanh(x_t)
	action = y_t * self.action_scale + self.action_bias
	log_prob = normal.log_prob(x_t)
	log_prob -= torch.log(self.action_scale * (1 - y_t.pow(2)) + 1e-6)
	log_prob = log_prob.sum(1, keepdim=True)
	mean = torch.tanh(mean) * self.action_scale + self.action_bias
	return action, log_prob, mean


	class CLIPVAE(nn.Module):
	"""Bottleneck autoencoder over CLIP caption embeddings (77 x 768 <-> 64)."""

	def __init__(self, latent_dim=64, seq_len=77, embed_dim=768, n_layers=4, n_heads=8, dropout=0.1, decoder_final_norm=True):
	super().__init__()
	self.seq_len = seq_len
	self.embed_dim = embed_dim

	self.pos_embed = nn.Parameter(torch.randn(1, seq_len, embed_dim) * 0.02)

	encoder_layer = nn.TransformerEncoderLayer(
	d_model=embed_dim, nhead=n_heads, dim_feedforward=2048,
	dropout=dropout, batch_first=True, norm_first=True
	)
	self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)

	self.bottleneck_norm = nn.LayerNorm(seq_len * embed_dim)

	self.fc_mu = nn.Linear(seq_len * embed_dim, latent_dim)
	self.fc_var = nn.Linear(seq_len * embed_dim, latent_dim)
	self.fc_decoder_input = nn.Linear(latent_dim, seq_len * embed_dim)

	decoder_layer = nn.TransformerEncoderLayer(
	d_model=embed_dim, nhead=n_heads, dim_feedforward=2048,
	dropout=dropout, batch_first=True, norm_first=True
	)
	self.transformer_decoder = nn.TransformerEncoder(decoder_layer, num_layers=n_layers)

	self.decoder_final_norm = nn.LayerNorm(embed_dim) if decoder_final_norm else None

	def reparameterize(self, mu, logvar):
	logvar = torch.clamp(logvar, min=-6, max=6)
	std = torch.exp(0.5 * logvar)
	eps = torch.randn_like(std)
	return mu + eps * std

	def encode(self, x):
	batch_size = x.shape[0]
	x = x + self.pos_embed
	x_encoded = self.transformer_encoder(x)

	x_flat = x_encoded.reshape(batch_size, -1)
	x_flat = self.bottleneck_norm(x_flat)

	mu = self.fc_mu(x_flat)
	logvar = self.fc_var(x_flat)
	return mu, logvar

	def decode(self, z):
	batch_size = z.shape[0]
	x_decode = self.fc_decoder_input(z).view(batch_size, self.seq_len, self.embed_dim)
	x_decode = x_decode + self.pos_embed

	out = self.transformer_decoder(x_decode)

	out = self.decoder_final_norm(out) if self.decoder_final_norm is not None else out
	return out

	def forward(self, x):
	mu, logvar = self.encode(x)
	z = self.reparameterize(mu, logvar)
	out = self.decode(z)
	return out, mu, logvar


	# ===========================================================================
	# Constants (paper configuration: SD v1.4 + DDIM, T = 50, g = 7.5)
	# ===========================================================================

	CKPT_REPO = "S-Karnik/rads-model-checkpoints"
	VAE_FILE = "clip_vae/step=45000.pth" # CLIP caption-embedding autoencoder (SD v1.4 / RealVis)
	TRAIN_SEEDS = [41, 42, 43, 44, 45] # the 5 SAC training seeds reported in the paper
	DEFAULT_TRAIN_SEED = 42
	CAPTION_DIM = 64
	TIME_EMBED_DIM = 16
	BASE_MODEL = "CompVis/stable-diffusion-v1-4"

	device = "cuda"


	def _load_state_dict(path: str):
	"""Load a torch state dict, preferring the safe weights-only path."""
	try:
	return torch.load(path, map_location="cpu", weights_only=True)
	except Exception:
	return torch.load(path, map_location="cpu", weights_only=False)


	def _load_rads(train_seed: int):
	"""Load the SAC encoder + actor checkpoints trained for `train_seed`."""
	prefix = f"RADS_SDV1_DDIM/checkpoints_SDV1_DDIM_seed_{train_seed}"
	encoder = SharedEncoder(CAPTION_DIM, time_embed_dim=TIME_EMBED_DIM, latent_resolution=64)
	encoder.load_state_dict(_load_state_dict(hf_hub_download(CKPT_REPO, f"{prefix}/sac_encoder_best.pth")))
	encoder.eval()

	actor = Actor(256, action_dim=CAPTION_DIM, action_space=None)
	actor.load_state_dict(_load_state_dict(hf_hub_download(CKPT_REPO, f"{prefix}/sac_actor_best.pth")))
	actor.eval()
	return encoder, actor


	# ===========================================================================
	# Model loading (module scope, eager, on the GPU — ZeroGPU defers the move)
	# ===========================================================================

	print("[RADS] loading Stable Diffusion v1.4 (fp16, DDIM)...")
	scheduler = DDIMScheduler.from_pretrained(BASE_MODEL, subfolder="scheduler")
	pipe = StableDiffusionPipeline.from_pretrained(
	BASE_MODEL,
	scheduler=scheduler,
	variant="fp16",
	torch_dtype=torch.float16,
	safety_checker=None,
	requires_safety_checker=False,
	)
	pipe.set_progress_bar_config(disable=True)
	pipe = pipe.to(device)

	print("[RADS] loading RADS steering components...")
	clip_vae = CLIPVAE(latent_dim=CAPTION_DIM, decoder_final_norm=False)
	clip_vae.load_state_dict(_load_state_dict(hf_hub_download(CKPT_REPO, VAE_FILE)))
	clip_vae.eval()
	clip_vae = clip_vae.to(device)

	agents = {}
	for train_seed in TRAIN_SEEDS:
	enc, act = _load_rads(train_seed)
	agents[train_seed] = (enc.to(device), act.to(device))
	print(f"[RADS] ready — loaded SAC agents for train seeds {TRAIN_SEEDS}")


	# ===========================================================================
	# Inference — replicates utils/run/generate.py (test_step, "sac" branch) from
	# the official repo 1:1: per-step SAC action in the 64-d caption latent space,
	# steered embedding replaces the conditional branch of classifier-free guidance.
	# ===========================================================================


	def _select_action(encoder, actor, latents, current_prompt_latents, t_norm):
	"""Deterministic (evaluate=True) SAC action for the current state."""
	with torch.no_grad():
	state_img = latents.to(device=device, dtype=torch.float32)
	state_cap = current_prompt_latents.to(device=device, dtype=torch.float32)
	if state_img.dim() == 3:
	state_img = state_img.unsqueeze(0)
	if state_cap.dim() == 2:
	state_cap = state_cap.unsqueeze(0)
	features = encoder(state_img, state_cap, t_norm)
	_, _, action = actor.sample(features) # evaluate=True -> deterministic mean
	return action.detach()


	def _generate_one(prompt, seed, num_steps, guidance, steer_scale, train_seed,
	progress=None, prog_range=(0.0, 1.0)):
	"""Run one SD v1.4 DDIM denoising trajectory; steer it with RADS if scale > 0.

	Returns a single PIL image.
	"""
	encoder, actor = agents[int(train_seed)]

	# --- prepare (mirrors wah's prepare_embeds / prepare_timesteps / prepare_latents)
	prompt_embeds, negative_embeds = pipe.encode_prompt(
	prompt=prompt,
	device=device,
	num_images_per_prompt=1,
	do_classifier_free_guidance=True,
	negative_prompt="",
	)
	embeds = torch.cat([negative_embeds, prompt_embeds], dim=0) # [2, 77, 768] fp16 (neg first)

	generator = torch.Generator(device=device).manual_seed(int(seed))
	pipe.scheduler.set_timesteps(num_steps, device=device)
	timesteps = pipe.scheduler.timesteps
	latents = pipe.prepare_latents(
	1, pipe.unet.config.in_channels, 512, 512,
	embeds.dtype, device, generator, None,
	)

	steer = steer_scale > 0.0
	if steer:
	embeds32 = embeds.to(dtype=torch.float32)
	with torch.no_grad():
	# caption latents fed to the policy: [B, 2, 64] = (neg, pos)
	prompt_mu, _ = clip_vae.encode(embeds32)
	neg_mu, pos_mu = prompt_mu.chunk(2)
	current_prompt_latents = torch.stack([neg_mu, pos_mu], dim=1)
	# mu of the conditional embedding (constant across steps; the reference
	# recomputes it every step from the same input, which is identical)
	cond_mu, _ = clip_vae.encode(embeds32.chunk(2)[1])
	neg_embeds = embeds.chunk(2)[0]

	with torch.no_grad():
	for i, t in enumerate(timesteps):
	if steer:
	action_latent = _select_action(
	encoder, actor, latents, current_prompt_latents, i / num_steps
	)
	z_steered = cond_mu + (action_latent * steer_scale)
	steered_cond_embedding = clip_vae.decode(z_steered)
	if steered_cond_embedding.dim() == 4:
	steered_cond_embedding = steered_cond_embedding.squeeze(1)
	steered_embeds = torch.cat(
	[neg_embeds, steered_cond_embedding.to(dtype=latents.dtype)], dim=0
	)
	step_embeds = steered_embeds
	else:
	step_embeds = embeds

	# --- predict_noise (classifier-free guidance: one batched forward pass)
	latent_model_input = torch.cat([latents] * 2)
	noise_pred = pipe.unet(
	latent_model_input, t, encoder_hidden_states=step_embeds, return_dict=False
	)[0]

	# --- perform_guidance
	noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
	noise_pred = noise_pred_uncond + guidance * (noise_pred_text - noise_pred_uncond)

	# --- step (DDIM, eta = 0)
	latents = pipe.scheduler.step(
	noise_pred, t, latents, generator=generator, return_dict=False
	)[0]

	if progress is not None:
	lo, hi = prog_range
	progress(lo + (i + 1) / num_steps * (hi - lo))

	# --- decode
	with torch.no_grad():
	decoded = pipe.vae.decode(
	latents / pipe.vae.config.scaling_factor, return_dict=False, generator=generator
	)[0]
	images = pipe.image_processor.postprocess(
	decoded, output_type="pil", do_denormalize=[True] * decoded.shape[0]
	)
	return images[0]


	@spaces.GPU(duration=20)
	def generate(prompt, steer_scale=1.0, train_seed=DEFAULT_TRAIN_SEED, seed=0,
	randomize_seed=False, num_steps=50, guidance_scale=7.5,
	progress=gr.Progress(track_tqdm=True)):
	"""Generate a baseline and a RADS-steered image for the same prompt and seed.

	Args:
	prompt: text prompt for Stable Diffusion v1.4.
	steer_scale: RADS steering strength (LATENT_SCALE in the paper's code);
	0 disables steering and reproduces the plain SD v1.4 baseline.
	train_seed: which of the 5 published SAC training seeds to use.
	seed: initial latent seed (shared by both images, so the pair differs
	only through steering).
	randomize_seed: draw a fresh random seed for this run.
	num_steps: number of DDIM denoising steps (paper default: 50).
	guidance_scale: classifier-free guidance scale (paper default: 7.5).

	Returns:
	(baseline image, RADS-steered image, seed used)
	"""
	if randomize_seed:
	seed = random.randint(0, 2**31 - 1)
	seed = int(seed)
	train_seed = int(train_seed)

	t0 = time.perf_counter()
	baseline = _generate_one(
	prompt, seed, int(num_steps), float(guidance_scale), 0.0, train_seed,
	progress=progress, prog_range=(0.0, 0.5),
	)
	t1 = time.perf_counter()
	steered = _generate_one(
	prompt, seed, int(num_steps), float(guidance_scale), float(steer_scale), train_seed,
	progress=progress, prog_range=(0.5, 1.0),
	)
	t2 = time.perf_counter()
	print(f"[RADS] prompt={prompt!r} seed={seed} train_seed={train_seed} "
	f"scale={steer_scale} baseline={t1 - t0:.2f}s steered={t2 - t1:.2f}s")
	return baseline, steered, seed


	# ===========================================================================
	# UI
	# ===========================================================================

	CSS = """
	#col-container { max-width: 1200px; margin: 0 auto; }
	.dark .gradio-container { color: var(--body-text-color); }
	"""

	EXAMPLES = [
	["satisfied happy scientist"],
	["fire us army garden flag single"],
	["The No Limits Business Woman Podcast"],
	["train that glorious with st robin western rug"],
	["gold colorful transit industry protected sign mosaic rug"],
	["Michael Fassbender to Star In Assassin's Creed Movie"],
	]

	with gr.Blocks() as demo:
	with gr.Column(elem_id="col-container"):
	gr.Markdown(
	"""
	# RADS: Steering Away from Memorization in Text-to-Image Diffusion

	RADS (Reachability-constrained reinforcement learning for Diffusion Steering) is an
	inference-time controller that steers Stable Diffusion v1.4 away from reproducing
	memorized training images. A SAC policy, trained offline against a reachability-based
	safety critic, predicts a small perturbation in a learned 64-d caption-embedding latent
	space at every denoising step — the diffusion backbone itself is never modified.

	This demo generates the same prompt and seed twice: plain SD v1.4 (left) vs.
	RADS-steered (right). Both images share the identical initial latent, so any
	difference comes purely from steering. Prompts below are memorization-prone prompts
	from the paper's evaluation sets (Webster 2023 / MemBench).

	[[Paper](https://huggingface.co/papers/2603.00140)] ·
	[Project page](https://s-karnik.github.io/rads-memorization-project-page) ·
	[Code](https://github.com/S-Karnik/rads-paper-repo) ·
	[Checkpoints](https://huggingface.co/S-Karnik/rads-model-checkpoints)
	"""
	)

	prompt = gr.Textbox(
	label="Prompt",
	placeholder="A memorization-prone prompt, e.g. 'satisfied happy scientist'",
	lines=2,
	)
	with gr.Row():
	steer_scale = gr.Slider(
	0.0, 2.0, value=1.0, step=0.05,
	label="Steering strength (0 = RADS off, 1 = paper default)",
	info="Scales the SAC action added to the caption-embedding latent (LATENT_SCALE).",
	)
	run_button = gr.Button("Generate", variant="primary", scale=1)

	with gr.Row():
	baseline_image = gr.Image(
	label="Baseline — SD v1.4 (no steering)", type="pil", height=420
	)
	steered_image = gr.Image(
	label="RADS-steered — same seed", type="pil", height=420
	)

	with gr.Accordion("Advanced settings", open=False):
	with gr.Row():
	train_seed = gr.Dropdown(
	choices=[str(s) for s in TRAIN_SEEDS],
	value=str(DEFAULT_TRAIN_SEED),
	label="SAC training seed",
	info="The paper reports 5 independently trained SAC agents (seeds 41–45).",
	)
	num_steps = gr.Slider(
	10, 100, value=50, step=1, label="DDIM steps (paper: 50)"
	)
	guidance_scale = gr.Slider(
	1.0, 15.0, value=7.5, step=0.5, label="Guidance scale (paper: 7.5)"
	)
	with gr.Row():
	seed = gr.Slider(
	0, 2**31 - 1, value=0, step=1, label="Seed"
	)
	randomize_seed = gr.Checkbox(
	value=False, label="Randomize seed (both images share it)"
	)

	gr.Examples(
	examples=EXAMPLES,
	inputs=[prompt],
	outputs=[baseline_image, steered_image, seed],
	fn=generate,
	cache_examples=True,
	cache_mode="lazy",
	)

	inputs = [prompt, steer_scale, train_seed, seed, randomize_seed, num_steps, guidance_scale]
	outputs = [baseline_image, steered_image, seed]
	run_button.click(fn=generate, inputs=inputs, outputs=outputs)

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