Spaces:

nvidia
/

nv-reason-ct

Running on Zero

filetype

3df4921 about 1 month ago

20.8 kB

"""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 = "<think>\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('<h1 style="text-align:center; margin: 0.2em 0; color: green;">NV-Reason-CT Demo. Check out the model card details <a href="https://huggingface.co/nvidia/NV-Reason-CT" target="_blank">here</a>.</h1>')

    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)

	"""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 = "<think>\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('<h1 style="text-align:center; margin: 0.2em 0; color: green;">NV-Reason-CT Demo. Check out the model card details <a href="https://huggingface.co/nvidia/NV-Reason-CT" target="_blank">here</a>.</h1>')

	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)