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)