app.py

"""
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)