app.py

from __future__ import annotations

import os
import threading
from dataclasses import dataclass

try:
    import spaces
except ImportError:
    class _LocalSpaces:
        @staticmethod
        def GPU(*args, **kwargs):
            del args, kwargs

            def decorator(fn):
                return fn

            return decorator

    spaces = _LocalSpaces()

import gradio as gr
import numpy as np
import torch
from diffusers import AutoPipelineForImage2Image
from PIL import Image, ImageOps

from stimulus_synthesis.generators.diffusers_t2i import DiffusersTextToImageAdapter
from stimulus_synthesis.media.normalize import videos_to_b_t_c_h_w
from stimulus_synthesis.neuro import available_rois, resolve_driving_voxels
from stimulus_synthesis.pipeline import NevoPipeline, _StaticImageToVideo


MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "mzx/NEvo")
MODEL_REVISION = os.getenv("MODEL_REVISION", "81ab95d6395b51620632e455ee5177759f74eaba")
IMAGE_MODEL_ID = os.getenv("IMAGE_MODEL_ID", "stabilityai/sdxl-turbo")
DEVICE = os.getenv("NEVO_DEVICE", "cuda")
IMAGE_SIZE = 512
SCORE_FRAMES = 16
ROI_NAMES = available_rois()
ROI_DESCRIPTIONS = {
    "EBA": "Bodies, body parts, poses, and bodily actions",
    "FFA": "Faces, facial configuration, viewpoint, and identity",
    "LOC": "Recognizable objects and object shape",
    "MT": "Visual motion, direction, and speed",
    "PPA": "Places, scenes, buildings, and spatial layout",
    "RSC": "Navigation, scene orientation, and familiar places",
    "V1": "Edges, contrast, orientation, and fine visual detail",
    "V2": "Contours, textures, boundaries, and local patterns",
    "V3": "Shape, depth, spatial structure, and dynamic form",
    "V4": "Color, curvature, texture, and complex visual form",
    "aSTS": "Higher-level person, social, and semantic information",
    "pSTS": "Biological motion, gaze, and social interaction",
}
ROI_CHOICES = [
    ("Auto - strongest predicted response", "Auto"),
    *[(f"{roi} - {ROI_DESCRIPTIONS[roi]}", roi) for roi in ROI_NAMES],
]


@dataclass
class Runtime:
    nevo: NevoPipeline


_runtime: Runtime | None = None
_runtime_lock = threading.Lock()
_inference_lock = threading.Lock()


def get_runtime() -> Runtime:
    global _runtime
    if _runtime is not None:
        return _runtime

    with _runtime_lock:
        if _runtime is not None:
            return _runtime
        if DEVICE == "cuda" and not torch.cuda.is_available():
            raise gr.Error("This Space requires CUDA GPU hardware.")
        if DEVICE == "mps" and not torch.backends.mps.is_available():
            raise gr.Error("MPS is not available in this environment.")
        if DEVICE not in {"cuda", "mps"}:
            raise gr.Error(f"Unsupported accelerator: {DEVICE}")

        image_pipe = AutoPipelineForImage2Image.from_pretrained(
            IMAGE_MODEL_ID,
            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
            use_safetensors=True,
        )
        image_pipe.enable_attention_slicing()
        image_generator = DiffusersTextToImageAdapter(
            IMAGE_MODEL_ID,
            device=DEVICE,
            pipeline=image_pipe,
        )
        nevo = NevoPipeline.from_pretrained(
            MODEL_REPO_ID,
            revision=MODEL_REVISION,
            device=DEVICE,
            text_to_image=image_generator,
        )
        nevo._ensure_components(device=DEVICE, image_only=True)
        _runtime = Runtime(nevo=nevo)
        return _runtime


def prepare_image(image: Image.Image | None) -> Image.Image:
    if image is None:
        raise gr.Error("Upload an image first.")
    return ImageOps.fit(
        image.convert("RGB"),
        (IMAGE_SIZE, IMAGE_SIZE),
        method=Image.Resampling.LANCZOS,
    )


def base_encoder_scorer(nevo: NevoPipeline):
    scorer = nevo.scorer
    while hasattr(scorer, "scorer"):
        scorer = scorer.scorer
    return scorer


def roi_profile(image: Image.Image, nevo: NevoPipeline) -> tuple[str, list[list[object]]]:
    scorer = base_encoder_scorer(nevo)
    static_video = _StaticImageToVideo(num_frames=SCORE_FRAMES).generate(image, "")
    batch = videos_to_b_t_c_h_w(
        [static_video],
        size=224,
        num_frames=SCORE_FRAMES,
    ).to(scorer.device)

    with torch.inference_mode():
        prediction = getattr(scorer.encoder, scorer.encoder_call)(batch)[0]
    prediction = prediction.detach().float().cpu().numpy()
    prediction = (prediction - prediction.mean()) / (prediction.std() + 1e-6)

    scores = []
    for roi in ROI_NAMES:
        mask = resolve_driving_voxels(roi)
        scores.append((roi, float(prediction[mask].mean())))
    scores.sort(key=lambda item: item[1], reverse=True)
    table = [
        [roi, ROI_DESCRIPTIONS[roi], round(score, 4)]
        for roi, score in scores
    ]
    return scores[0][0], table


@spaces.GPU(duration=300, size="large")
def analyze_image(image: Image.Image | None):
    source = prepare_image(image)
    with _inference_lock:
        inferred_roi, table = roi_profile(source, get_runtime().nevo)
    return inferred_roi, table


def generation_duration(image, roi_choice, strength, evaluations, seed):
    del image, roi_choice, strength, seed
    return max(300, 120 + 12 * max(2, int(evaluations)))


@spaces.GPU(duration=generation_duration, size="large")
def generate_image(
    image: Image.Image | None,
    roi_choice: str,
    strength: float,
    evaluations: int,
    seed: int,
):
    source = prepare_image(image)
    evaluations = max(2, int(evaluations))
    seed = int(seed)

    with _inference_lock:
        nevo = get_runtime().nevo
        inferred_roi, table = roi_profile(source, nevo)
        target_roi = inferred_roi if roi_choice == "Auto" else roi_choice
        result = nevo(
            roi=target_roi,
            image_only=True,
            progress=False,
            image_max_evals=evaluations,
            population_size=min(4, evaluations),
            image_batch_size=1,
            image_kwargs={
                "image": source,
                "strength": float(strength),
                "num_inference_steps": 4,
                "guidance_scale": 0.0,
                "height": IMAGE_SIZE,
                "width": IMAGE_SIZE,
            },
            score_size=224,
            num_frames=SCORE_FRAMES,
            seed=seed,
        )

    return (
        result.best.image,
        inferred_roi,
        target_roi,
        result.best_prompt,
        round(float(result.best_score), 5),
        table,
    )


with gr.Blocks(title="NEvo Image-to-Image - Community Adaptation") as demo:
    gr.Markdown(
        "# NEvo Image-to-Image\n"
        "Independent community adaptation of "
        "[NEvo by EPFL NeuroAI](https://huggingface.co/epfl-neuroai/NEvo). "
        "We are not authors of the original NEvo paper."
    )
    with gr.Row():
        with gr.Column(scale=1):
            source_image = gr.Image(type="pil", label="Source image", height=420)
            roi_choice = gr.Dropdown(
                choices=ROI_CHOICES,
                value="Auto",
                label="Target ROI",
            )
            strength = gr.Slider(
                minimum=0.2,
                maximum=1.0,
                value=0.65,
                step=0.05,
                label="Edit strength",
            )
            evaluations = gr.Number(
                minimum=2,
                value=8,
                precision=0,
                label="Search evaluations",
            )
            seed = gr.Number(value=0, precision=0, label="Seed")
            with gr.Row():
                analyze_button = gr.Button("Analyze", variant="secondary")
                generate_button = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            result_image = gr.Image(type="pil", label="Optimized image", height=420)
            with gr.Row():
                inferred_roi = gr.Textbox(label="Inferred ROI", interactive=False)
                target_roi = gr.Textbox(label="Optimized ROI", interactive=False)
                best_score = gr.Number(label="Best score", interactive=False)
            best_prompt = gr.Textbox(label="Evolved prompt", interactive=False)
            profile = gr.Dataframe(
                headers=["ROI", "Corresponds to", "Relative predicted response"],
                datatype=["str", "str", "number"],
                interactive=False,
            )

    analyze_button.click(
        fn=analyze_image,
        inputs=[source_image],
        outputs=[inferred_roi, profile],
    )
    generate_button.click(
        fn=generate_image,
        inputs=[source_image, roi_choice, strength, evaluations, seed],
        outputs=[
            result_image,
            inferred_roi,
            target_roi,
            best_prompt,
            best_score,
            profile,
        ],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=8).launch()