"""NV-Reason-CT Space app.""" import contextlib import io import os import re import shutil import tempfile import types import uuid from collections import OrderedDict from pathlib import Path from threading import Lock, Thread SPACES_DIR = Path(__file__).resolve().parent import gradio as gr import numpy as np import spaces import torch from PIL import Image from transformers import AutoModelForImageTextToText, AutoProcessor, TextIteratorStreamer DEFAULT_MODEL = "nvidia/NV-Reason-CT" EXAMPLES_DIR = SPACES_DIR / "examples" pretrained_model_name_or_path = os.environ.get("MODEL", DEFAULT_MODEL) GRADIO_SERVER_PORT = int(os.environ.get("GRADIO_SERVER_PORT", "7860")) if not torch.cuda.is_available(): raise RuntimeError("CUDA is not available.") DEVICE = "cuda" MODEL_DTYPE = torch.bfloat16 ATTN_IMPLEMENTATION = str(os.environ.get("ATTN_IMPLEMENTATION", "sdpa")) DEFAULT_PROMPT = "write a structured chest CT report" REGION_CHOICES = ["Chest", "Abdomen"] DEFAULT_REGION = "Chest" REGION_TO_ANATOMY = {"Chest": "chest", "Abdomen": "abdomen"} STRUCTURED_REPORT_PROMPTS = { "Chest": "write a structured chest CT report", "Abdomen": "write a structured abdominal CT report", } REASONING_ANALYSIS_PROMPTS = { "Chest": "full chest CT reasoning analysis", "Abdomen": "full abdominal CT reasoning analysis", } MAX_CACHED_VOLUMES = int(os.environ.get("VLM3D_APP_MAX_CACHED_VOLUMES", "10")) VOLUME_CACHE = OrderedDict() VOLUME_CACHE_LOCK = Lock() def structured_report_prompt(region): return STRUCTURED_REPORT_PROMPTS.get(region, DEFAULT_PROMPT) def reasoning_analysis_prompt(region): return REASONING_ANALYSIS_PROMPTS.get(region, REASONING_ANALYSIS_PROMPTS[DEFAULT_REGION]) def _cleanup_case(case): if not case: return managed_path = case.get("managed_path") if managed_path: Path(managed_path).unlink(missing_ok=True) def _cache_case(case): key = uuid.uuid4().hex with VOLUME_CACHE_LOCK: VOLUME_CACHE[key] = case VOLUME_CACHE.move_to_end(key) while len(VOLUME_CACHE) > MAX_CACHED_VOLUMES: _, old_case = VOLUME_CACHE.popitem(last=False) _cleanup_case(old_case) return key def _drop_cached_volume(key): if key is None: return with VOLUME_CACHE_LOCK: case = VOLUME_CACHE.pop(key, None) _cleanup_case(case) def _get_cached_case(key): if key is None: return None with VOLUME_CACHE_LOCK: case = VOLUME_CACHE.get(key) if case is not None: VOLUME_CACHE.move_to_end(key) if case is None: raise gr.Error("Cached volume expired. Please upload/select the image again.") return case def _load_model(): load_kwargs = dict( pretrained_model_name_or_path=pretrained_model_name_or_path, dtype=MODEL_DTYPE, attn_implementation=ATTN_IMPLEMENTATION, trust_remote_code=True, ) with contextlib.redirect_stdout(io.StringIO()): loaded = AutoModelForImageTextToText.from_pretrained(**load_kwargs) if not (hasattr(loaded, "model") and hasattr(loaded.model, "vision3d")): raise RuntimeError( f"AutoModelForImageTextToText resolved to {type(loaded).__name__}, " "which has no vision3d tower. Check the model repo auto_map." ) return loaded.eval().to(DEVICE) model = _load_model() with contextlib.redirect_stdout(io.StringIO()): processor = AutoProcessor.from_pretrained( pretrained_model_name_or_path=pretrained_model_name_or_path, trust_remote_code=True, ) processor.tokenizer.padding_side = "left" if getattr(processor, "image_processor_3d", None) is None: raise RuntimeError("Processor has no image_processor_3d -- not the right checkpoint.") def _get_generation_normalize_mode(): return getattr(processor.image_processor_3d, "normalize_mode", None) def _prepare_volume_for_generation(vol): image = torch.as_tensor(vol).unsqueeze(0) normalize_mode = _get_generation_normalize_mode() return processor.image_processor_3d._apply_normalization(image, normalize_mode) def identity_image3d_preprocess(self, images, *, normalize_mode=None, **kwargs): pixel_values = [] image_grid_thws = [] for image in images: if isinstance(image, np.ndarray): image = torch.from_numpy(image) if not isinstance(image, torch.Tensor): raise TypeError(f"images3d items must be tensor/ndarray, got {type(image)!r}") if image.ndim == 3: image = image.unsqueeze(0) pixel_values.append(image.contiguous()) image_grid_thws.append(torch.tensor(self.final_grid_size)) if len(pixel_values) == 1: pixel_values = pixel_values[0].unsqueeze(0) else: pixel_values = torch.stack(pixel_values, dim=0) return { "pixel_values": pixel_values, "image_grid_thw": torch.stack(image_grid_thws, dim=0), } processor.image_processor_3d.preprocess = types.MethodType( identity_image3d_preprocess, processor.image_processor_3d, ) EXAMPLE_CASES = [ {"path": str(EXAMPLES_DIR / "example_1.nii.gz"), "prompt": DEFAULT_PROMPT}, {"path": str(EXAMPLES_DIR / "example_2.nii.gz"), "prompt": DEFAULT_PROMPT}, {"path": str(EXAMPLES_DIR / "example_3.nii.gz"), "prompt": DEFAULT_PROMPT}, ] def _as_bool(value): if isinstance(value, str): return value == "True" return bool(value) _SPECIAL_TOKEN_RE = re.compile(r"<\|[^|]*\|>") def _strip_special_tokens(text): return _SPECIAL_TOKEN_RE.sub("", text).strip() def _history_content_to_text(content): if content is None: return "" if isinstance(content, str): return _strip_special_tokens(content) if isinstance(content, dict): if isinstance(content.get("text"), str): return _strip_special_tokens(content["text"]) return "" if isinstance(content, list): parts = [] for item in content: if isinstance(item, str): parts.append(item) elif isinstance(item, dict) and isinstance(item.get("text"), str): parts.append(item["text"]) return _strip_special_tokens(" ".join(p.strip() for p in parts if p)) return "" def _prepare_inputs(text, history, volume_key, region, enable_thinking): if volume_key is None: raise gr.Error("Please upload a .nii/.nii.gz file.", duration=3, print_exception=False) if text is None or len(text.strip()) == 0: raise gr.Error("Please input a query.", duration=3, print_exception=False) case = _get_cached_case(volume_key) vol = _get_case_crop(case, region) image = _prepare_volume_for_generation(vol) messages = [] for h in history: role = h.get("role") if role not in ("user", "assistant"): continue text_str = _history_content_to_text(h.get("content")) if not text_str: continue messages.append({"role": role, "content": [{"type": "text", "text": text_str}]}) messages.append({"role": "user", "content": [{"type": "text", "text": text.strip()}]}) messages[0]["content"].insert(0, {"type": "image"}) enable_thinking_b = _as_bool(enable_thinking) prompt = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking_b, ) return processor(text=prompt, images3d=[image], return_tensors="pt") @spaces.GPU def _generate_on_gpu(inputs, temperature, repetition_penalty, do_sample): inputs = inputs.to(DEVICE) streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True) generation_args = dict(inputs, streamer=streamer, max_new_tokens=4096, use_cache=True) generation_args["repetition_penalty"] = float(repetition_penalty) generation_args["do_sample"] = _as_bool(do_sample) if generation_args["do_sample"]: generation_args["temperature"] = float(temperature) with torch.inference_mode(): thread = Thread(target=model.generate, kwargs=generation_args) thread.start() yield "..." buffer = "" for new_text in streamer: buffer += new_text yield buffer def model_inference(text, history, volume_key, region, temperature, repetition_penalty, do_sample, enable_thinking): inputs = _prepare_inputs(text, history, volume_key, region, enable_thinking) prefix = "\n" if _as_bool(enable_thinking) else "" for chunk in _generate_on_gpu(inputs, temperature, repetition_penalty, do_sample): yield prefix + chunk def _window_to_uint8(slice2d: np.ndarray) -> np.ndarray: x = np.asarray(slice2d, dtype=np.float32) x = np.clip(x, -500, 500) x = (x + 500) / 1000.0 return (x * 255.0).astype(np.uint8) def _is_nifti(path: str) -> bool: p = path.lower() return p.endswith(".nii") or p.endswith(".nii.gz") def _copy_to_managed_temp(path: str) -> str: suffix = ".nii.gz" if path.lower().endswith(".nii.gz") else ".nii" tmp = tempfile.NamedTemporaryFile(prefix="vlm3d_upload_", suffix=suffix, delete=False) tmp_path = tmp.name tmp.close() try: shutil.copyfile(path, tmp_path) except Exception: Path(tmp_path).unlink(missing_ok=True) raise return tmp_path def _make_case(path: str, *, managed_copy: bool): if not _is_nifti(path): raise gr.Error("Please upload a .nii or .nii.gz file.") managed_path = _copy_to_managed_temp(path) if managed_copy else None return { "path": managed_path or path, "managed_path": managed_path, "crops": {}, "lock": Lock(), } def _load_region_crop(path: str, anatomy_region: str) -> np.ndarray: image = processor.image_processor_3d.load_image(path, normalize_mode=0, anatomy_region=anatomy_region) if isinstance(image, torch.Tensor): vol = image.detach().cpu().numpy().squeeze() else: vol = np.asarray(image).squeeze() return vol def _get_case_crop(case, region) -> np.ndarray: anatomy_region = REGION_TO_ANATOMY[region] with case["lock"]: if anatomy_region not in case["crops"]: case["crops"][anatomy_region] = _load_region_crop(case["path"], anatomy_region) return case["crops"][anatomy_region] def _make_preview_image(vol: np.ndarray, k: int) -> Image.Image: sl = vol[k, :, :] pil = Image.fromarray(_window_to_uint8(sl)) return pil.resize((384, 384), resample=Image.Resampling.BILINEAR) def _preview_update_for_case(case, region): vol = _get_case_crop(case, region) z = vol.shape[0] k = z // 2 return ( gr.update(value=_make_preview_image(vol, k), visible=True), gr.update(minimum=0, maximum=z - 1, value=k, visible=True), ) def _build_example_gallery(): rows = [] for case in EXAMPLE_CASES: case_obj = _make_case(case["path"], managed_copy=False) vol = _get_case_crop(case_obj, DEFAULT_REGION) preview = _make_preview_image(vol, vol.shape[0] // 2) rows.append([preview, case["prompt"]]) return rows EXAMPLE_ROWS = _build_example_gallery() def on_upload(file_obj): if file_obj is None: return ( gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), None, gr.update(minimum=0, maximum=191, value=96), gr.update(value=DEFAULT_REGION), ) case = _make_case(file_obj.name, managed_copy=True) try: viewer_update, slider_update = _preview_update_for_case(case, DEFAULT_REGION) except Exception: _cleanup_case(case) raise volume_key = _cache_case(case) return ( viewer_update, gr.update(visible=False), gr.update(visible=True), volume_key, slider_update, gr.update(value=DEFAULT_REGION), ) def on_slice_change(k, volume_key, region): if volume_key is None: return gr.update() vol = _get_case_crop(_get_cached_case(volume_key), region) z = vol.shape[0] k = int(np.clip(int(k), 0, z - 1)) pil = Image.fromarray(_window_to_uint8(vol[k, :, :])) return pil.resize((384, 384), resample=Image.Resampling.BILINEAR) def on_region_change(region, volume_key): report_prompt = structured_report_prompt(region) reasoning_prompt = reasoning_analysis_prompt(region) if volume_key is None: return gr.update(), gr.update(), [], [], [], None, report_prompt, report_prompt, reasoning_prompt viewer_update, slider_update = _preview_update_for_case(_get_cached_case(volume_key), region) return viewer_update, slider_update, [], [], [], None, report_prompt, report_prompt, reasoning_prompt def on_select_example(example_tuple, volume_key): _drop_cached_volume(volume_key) if example_tuple is None: return gr.update(), None, gr.update(), gr.update(), gr.update(), gr.update(), [], [], [], None example_id, _example_values = example_tuple case_info = EXAMPLE_CASES[int(example_id)] case = _make_case(case_info["path"], managed_copy=False) viewer_update, slider_update = _preview_update_for_case(case, DEFAULT_REGION) volume_key = _cache_case(case) return ( viewer_update, volume_key, gr.update(visible=True), gr.update(visible=False), slider_update, gr.update(value=DEFAULT_REGION), [], [], [], None, ) def on_change_file(volume_key): _drop_cached_volume(volume_key) return ( gr.update(value=None, visible=True), gr.update(value=None, visible=False), gr.update(visible=False), None, gr.update(minimum=0, maximum=191, value=96), gr.update(value=DEFAULT_REGION), [], [], [], None, DEFAULT_PROMPT, ) with gr.Blocks() as demo: gr.HTML('

NV-Reason-CT Demo. Check out the model card details here.

') textbox = gr.Textbox( show_label=False, placeholder="Enter your text here and press ENTER", render=False, submit_btn="Send", value=DEFAULT_PROMPT, ) with gr.Row(): with gr.Column(scale=1): nifti = gr.File( label="Upload .nii/.nii.gz", file_count="single", #file_types=[".nii", ".nii.gz"], visible=True, ) viewer = gr.Image( type="pil", label="Input crop", show_label=True, height=384, visible=False, ) slider = gr.Slider(minimum=0, maximum=191, value=96, step=1, label="Slice", visible=True) region_dropdown = gr.Dropdown( choices=REGION_CHOICES, value=DEFAULT_REGION, label="Region", ) change_btn = gr.Button("Change file", visible=False) state = gr.State(None, delete_callback=_drop_cached_volume) example_preview_input = gr.Image(type="pil", visible=False) nifti.change( fn=on_upload, inputs=nifti, outputs=[viewer, nifti, change_btn, state, slider, region_dropdown], queue=False, api_visibility="private", ) slider.change( fn=on_slice_change, inputs=[slider, state, region_dropdown], outputs=viewer, queue=False, api_visibility="private", ) with gr.Accordion("Generation Options", open=False): with gr.Row(): control_temperature = gr.Number( minimum=0.0, maximum=1.0, value=0.000001, step=0.000001, label="Temperature", ) control_repetition_penalty = gr.Dropdown( choices=[1.0, 1.05, 1.1, 1.2], value=1.0, label="Repetition Penalty", allow_custom_value=True, ) control_do_sample = gr.Dropdown( choices=["True", "False"], value="False", label="Do Sample", ) control_enable_thinking = gr.Checkbox(value=True, label="Enable Thinking") with gr.Accordion("Examples", open=True): ex = gr.Examples( examples=EXAMPLE_ROWS, inputs=[example_preview_input, textbox], cache_examples=False, label=None, ) with gr.Column(scale=2): chat_interface = gr.ChatInterface( fn=model_inference, chatbot=gr.Chatbot( render_markdown=True, sanitize_html=False, allow_tags=True, height="35vw", container=False, ), textbox=textbox, multimodal=False, fill_height=False, api_visibility="private", concurrency_limit=1, additional_inputs=[ state, region_dropdown, control_temperature, control_repetition_penalty, control_do_sample, control_enable_thinking, ], ) with gr.Row(): structured_report_btn = gr.Button( structured_report_prompt(DEFAULT_REGION), size="sm", scale=0, min_width=240 ) reasoning_analysis_btn = gr.Button( reasoning_analysis_prompt(DEFAULT_REGION), size="sm", scale=0, min_width=240 ) structured_report_btn.click( fn=structured_report_prompt, inputs=region_dropdown, outputs=textbox, queue=False, api_visibility="private", ) reasoning_analysis_btn.click( fn=reasoning_analysis_prompt, inputs=region_dropdown, outputs=textbox, queue=False, api_visibility="private", ) region_dropdown.change( fn=on_region_change, inputs=[region_dropdown, state], outputs=[ viewer, slider, chat_interface.chatbot, chat_interface.chatbot_state, chat_interface.chatbot_value, chat_interface.saved_input, textbox, structured_report_btn, reasoning_analysis_btn, ], queue=False, api_visibility="private", ) change_btn.click( fn=on_change_file, inputs=state, outputs=[ nifti, viewer, change_btn, state, slider, region_dropdown, chat_interface.chatbot, chat_interface.chatbot_state, chat_interface.chatbot_value, chat_interface.saved_input, textbox, ], queue=False, api_visibility="private", ) ex.load_input_event.then( fn=on_select_example, inputs=[ex.dataset, state], outputs=[ viewer, state, change_btn, nifti, slider, region_dropdown, chat_interface.chatbot, chat_interface.chatbot_state, chat_interface.chatbot_value, chat_interface.saved_input, ], queue=False, api_visibility="private", ) demo.queue(max_size=10) demo.launch(debug=False, server_name="0.0.0.0", server_port=GRADIO_SERVER_PORT)