""" PolypSteer / MedSteer — Counterfactual Endoscopic Synthesis via Training-Free Activation Steering. This Space loads a PixArt-α (512×512) pipeline LoRA-fine-tuned on the Kvasir endoscopy dataset (phamtrongthang/medsteer) and applies training-free activation steering to the cross-attention output of every DiT transformer block, producing a baseline image and a steered (concept-suppressed) counterfactual side-by-side. """ import os from copy import deepcopy os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # MUST come before torch / any CUDA import import torch import numpy as np import gradio as gr from peft import PeftModel from huggingface_hub import snapshot_download from diffusers import PixArtAlphaPipeline # ── Model identifiers ────────────────────────────────────────────────────── BASE_MODEL = "PixArt-alpha/PixArt-XL-2-512x512" LORA_REPO = "phamtrongthang/medsteer" DTYPE = torch.float16 # use fp16 weights — much smaller download # Concept pair for precomputed direction vectors POS_CONCEPT = "dyed lifted polyps" NEG_CONCEPT = "normal cecum" PROMPT_PREFIX = "An endoscopic image of " # ── Compatibility shim ────────────────────────────────────────────────────── # newer transformers removed FLAX_WEIGHTS_NAME; patch it back before diffusers import transformers.utils as _tu if not hasattr(_tu, "FLAX_WEIGHTS_NAME"): _tu.FLAX_WEIGHTS_NAME = "diffusion_flax_model.msgpack" # ── Model loading (module scope, no GPU needed — ZeroGPU hijack handles .to("cuda")) ── def load_pipeline() -> PixArtAlphaPipeline: """Load PixArt-α with LoRA adapters from phamtrongthang/medsteer.""" lora_path = snapshot_download(repo_id=LORA_REPO) pipe = PixArtAlphaPipeline.from_pretrained( BASE_MODEL, torch_dtype=DTYPE, variant="fp16", ) # Keep VAE in fp16 to match the rest of the pipeline (fp32 VAE causes # dtype mismatch with fp16 latents from the transformer). # Load LoRA adapters for transformer and text encoder. # ZeroGPU patches torch.cuda.is_available() to True at module scope, so # peft's infer_device() returns "cuda" and safe_load_file tries to use # CUDA — which fails because there's no GPU at startup. Temporarily # make CUDA "unavailable" so peft loads weights on CPU; the subsequent # pipe.to("cuda") is intercepted by ZeroGPU's hijack as expected. _orig_is_available = torch.cuda.is_available torch.cuda.is_available = lambda: False try: pipe.transformer = PeftModel.from_pretrained( pipe.transformer, os.path.join(lora_path, "transformer_lora"), is_trainable=False, torch_dtype=DTYPE, ) pipe.text_encoder = PeftModel.from_pretrained( pipe.text_encoder, os.path.join(lora_path, "text_encoder_lora"), is_trainable=False, torch_dtype=DTYPE, ) finally: torch.cuda.is_available = _orig_is_available pipe.to("cuda") return pipe pipe = load_pipeline() print("[PolypSteer] Pipeline loaded.") # ── Activation steering core ──────────────────────────────────────────────── class CrossAttentionHook: """Register a forward hook on every attn2 module to intercept cross-attention output. Modes: - "record": collect mean activation per step/block (no modification) - "suppress": subtract the aligned component along the direction vector - "baseline": no hook action (passthrough) """ def __init__(self): self.handles = [] self.mode = "baseline" self.direction_vectors = None self.suppress_scale = 2.0 self._current_step = 0 self._total_blocks = 0 self._current_block = 0 self._step_buffer = {"blocks": []} self._activation_cache = {} def attach(self, transformer): self.handles = [] self._total_blocks = 0 for i, block in enumerate(transformer.transformer_blocks): handle = block.attn2.register_forward_hook(self._make_hook(i)) self.handles.append(handle) self._total_blocks += 1 print(f"[PolypSteer] Attached hooks to {self._total_blocks} blocks.") def reset_state(self): self._current_step = 0 self._current_block = 0 self._step_buffer = {"blocks": []} self._activation_cache = {} def _make_hook(self, block_idx): def hook(module, input, output): # output is a tuple; the first element is the attention output if isinstance(output, tuple): activation = output[0] else: activation = output if self.mode == "suppress" and self.direction_vectors is not None: max_step = max(self.direction_vectors.keys()) num_step = ( self._current_step if self._current_step in self.direction_vectors else max_step ) if num_step > max_step: num_step = max_step if num_step in self.direction_vectors: blocks = self.direction_vectors[num_step].get("blocks", []) if block_idx < len(blocks): dv = torch.tensor( blocks[block_idx], device=activation.device, dtype=activation.dtype ).view(1, 1, -1) norm = torch.norm(activation, dim=2, keepdim=True) sim = torch.tensordot( activation, dv, dims=([2], [2]) ).view(activation.size(0), activation.size(1), 1) sim = torch.where(sim > 0, sim, torch.zeros_like(sim)) activation = activation - ( self.suppress_scale * sim ) * dv.expand(activation.size(0), activation.size(1), -1) activation = activation / ( torch.norm(activation, dim=2, keepdim=True) + 1e-8 ) activation = activation * norm # Record activations (always - matches original code) if activation.shape[0] > 1: captured = ( activation.detach().cpu().numpy()[len(activation) // 2:] .mean(axis=0).mean(axis=0) ) else: captured = activation.detach().cpu().numpy().mean(axis=0).mean(axis=0) self._step_buffer["blocks"].append(captured) # Track step/block progression self._current_block += 1 if self._current_block == self._total_blocks: self._current_block = 0 self._activation_cache[self._current_step] = self._step_buffer self._step_buffer = {"blocks": []} self._current_step += 1 if isinstance(output, tuple): return (activation,) + output[1:] return activation return hook steer_hook = CrossAttentionHook() steer_hook.attach(pipe.transformer) print("[PolypSteer] Hooks attached.") # ── Direction vector computation (runs inside @spaces.GPU on first call) ─── _direction_vectors_cache = None @torch.no_grad() def _compute_direction_vectors( pos_prompt: str, neg_prompt: str, num_images: int = 3, num_steps: int = 20, base_seed: int = 1000, ): """Capture activations for two concept prompts and compute mean-difference direction vectors. Returns a dict indexed as direction_vectors[step]["blocks"][block_idx]. Must be called inside @spaces.GPU — requires a real GPU. """ pos_activations = [] neg_activations = [] for label, prompt_text in [("pos", pos_prompt), ("neg", neg_prompt)]: for i in range(num_images): steer_hook.reset_state() steer_hook.mode = "record" seed = base_seed + i generator = torch.Generator(device="cuda").manual_seed(seed) pipe( prompt=prompt_text, num_inference_steps=num_steps, generator=generator, use_resolution_binning=False, ) cache = deepcopy(steer_hook._activation_cache) if label == "pos": pos_activations.append(cache) else: neg_activations.append(cache) # Compute direction vectors num_steps_actual = len(pos_activations[0]) direction_vectors = {} for step in range(num_steps_actual): direction_vectors[step] = {"blocks": []} num_blocks = len(pos_activations[0][step]["blocks"]) for block_idx in range(num_blocks): pos_layer = [ pos_activations[i][step]["blocks"][block_idx] for i in range(len(pos_activations)) ] pos_avg = np.mean(pos_layer, axis=0) neg_layer = [ neg_activations[i][step]["blocks"][block_idx] for i in range(len(neg_activations)) ] neg_avg = np.mean(neg_layer, axis=0) direction = pos_avg - neg_avg norm = np.linalg.norm(direction) if norm > 1e-8: direction = direction / norm direction_vectors[step]["blocks"].append(direction) return direction_vectors # ── Inference ─────────────────────────────────────────────────────────────── @spaces.GPU(duration=120) def generate( prompt: str, seed: int = 42, num_steps: int = 20, suppress_scale: float = 2.0, progress=gr.Progress(track_tqdm=True), ): """Generate a baseline endoscopic image and its steered counterfactual. The baseline image is produced from the fine-tuned PixArt-α model. The steered image suppresses concept-specific features (e.g. polyp appearance) via activation steering, showing what the same scene would look like without the pathological finding. Args: prompt: Text prompt describing the endoscopic scene. seed: RNG seed for reproducibility. num_steps: Number of denoising steps (20 is a good default). suppress_scale: Steering strength (1-3 work well; higher = more suppression). """ global _direction_vectors_cache seed = int(seed) num_steps = int(num_steps) # Compute direction vectors on first call (requires GPU) if _direction_vectors_cache is None: print("[PolypSteer] Computing direction vectors (first call)…") _direction_vectors_cache = _compute_direction_vectors( pos_prompt=f"{PROMPT_PREFIX}{POS_CONCEPT}", neg_prompt=f"{PROMPT_PREFIX}{NEG_CONCEPT}", num_images=3, num_steps=20, base_seed=1000, ) print("[PolypSteer] Direction vectors ready.") # ── Baseline ── steer_hook.reset_state() steer_hook.mode = "baseline" generator = torch.Generator(device="cuda").manual_seed(seed) baseline_img = pipe( prompt=prompt, num_inference_steps=num_steps, generator=generator, use_resolution_binning=False, ).images[0] # ── Steered (suppress) ── steer_hook.reset_state() steer_hook.mode = "suppress" steer_hook.direction_vectors = _direction_vectors_cache steer_hook.suppress_scale = suppress_scale generator = torch.Generator(device="cuda").manual_seed(seed) steered_img = pipe( prompt=prompt, num_inference_steps=num_steps, generator=generator, use_resolution_binning=False, ).images[0] # Reset hook state steer_hook.mode = "baseline" steer_hook.reset_state() return baseline_img, steered_img # ── Gradio UI ─────────────────────────────────────────────────────────────── CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# PolypSteer: Counterfactual Endoscopic Synthesis\n" "Training-free activation steering for endoscopic image generation. " "Generate a baseline endoscopic image and its steered counterfactual " "(with pathological features suppressed) using a PixArt-α model " "fine-tuned on Kvasir." ) with gr.Row(): prompt = gr.Textbox( label="Prompt", value=f"{PROMPT_PREFIX}{POS_CONCEPT}", show_label=True, container=False, scale=4, ) run_btn = gr.Button("Generate", variant="primary", scale=1) with gr.Row(): baseline_out = gr.Image( label="Baseline (fine-tuned model)", type="pil", height=512, ) steered_out = gr.Image( label="Steered (concept suppressed)", type="pil", height=512, ) with gr.Accordion("Advanced settings", open=False): seed = gr.Number(label="Seed", value=42, precision=0) num_steps = gr.Slider( label="Denoising steps", minimum=5, maximum=50, value=20, step=1, ) suppress_scale = gr.Slider( label="Suppress scale (steering strength)", minimum=0.0, maximum=5.0, value=2.0, step=0.1, ) gr.Examples( examples=[ [f"{PROMPT_PREFIX}{POS_CONCEPT}", 42, 20, 2.0], [f"{PROMPT_PREFIX}polyps", 42, 20, 2.0], [f"{PROMPT_PREFIX}ulcerative colitis", 42, 20, 2.0], [f"{PROMPT_PREFIX}dyed resection margins", 42, 20, 2.0], ], inputs=[prompt, seed, num_steps, suppress_scale], outputs=[baseline_out, steered_out], fn=generate, cache_examples=True, cache_mode="lazy", ) gr.Markdown( "**Model:** [PixArt-α](https://huggingface.co/PixArt-alpha/PixArt-XL-2-512x512) " "with LoRA adapters from [phamtrongthang/medsteer](https://huggingface.co/phamtrongthang/medsteer). \n" "**Paper:** [PolypSteer: Counterfactual Endoscopic Synthesis via " "Training-Free Activation Steering](https://huggingface.co/papers/2603.07066) \n" "**Code:** [GitHub](https://github.com/UARK-AICV/PolypSteer)" ) run_btn.click( fn=generate, inputs=[prompt, seed, num_steps, suppress_scale], outputs=[baseline_out, steered_out], ) demo.launch(mcp_server=True)