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