Upload 2 files

app.py

@@ -0,0 +1,146 @@
+import os
+import numpy as np
+import torch
+import gradio as gr
+from PIL import Image
+
+from diffusers import StableDiffusion3ControlNetPipeline
+from transformers import T5Tokenizer
+
+
+# --- Semantic palette from the model card ---
+CLASS_DICT_ATLAS = {
+    0: (0, 0, 0),
+    1: (255, 60, 0),
+    2: (255, 60, 232),
+    3: (134, 79, 117),
+    4: (125, 0, 190),
+    5: (117, 200, 191),
+    6: (230, 91, 101),
+    7: (255, 0, 155),
+    8: (75, 205, 155),
+    9: (100, 37, 200),
+}
+
+NAME_CLASS_DICT = {
+    0: "background",
+    1: "aorta",
+    2: "kidney_left",
+    3: "liver",
+    4: "postcava",
+    5: "stomach",
+    6: "gall_bladder",
+    7: "kidney_right",
+    8: "pancreas",
+    9: "spleen",
+}
+
+
+def mask_to_labels(mask_rgb: Image.Image) -> np.ndarray:
+    """Convert RGB mask -> label map using exact color matches."""
+    arr = np.asarray(mask_rgb.convert("RGB"), dtype=np.uint8)
+    h, w, _ = arr.shape
+    labels = np.zeros((h, w), dtype=np.int16)
+
+    for cls_id, rgb in CLASS_DICT_ATLAS.items():
+        m = np.all(arr == np.array(rgb, dtype=np.uint8), axis=-1)
+        labels[m] = cls_id
+    return labels
+
+
+# --- Load pipeline once (cached) ---
+MODEL_ID = "onkarsus13/Semantic-Control-Stable-diffusion-3-M-Mask2CT-Atlas"
+PIPE = None
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def load_pipe():
+    global PIPE
+    if PIPE is not None:
+        return PIPE
+
+    dtype = torch.float16 if DEVICE == "cuda" else torch.float32
+
+    pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
+        MODEL_ID,
+        torch_dtype=dtype,
+        safety_checker=None,
+        feature_extractor=None,
+    )
+
+    # Model card explicitly loads tokenizer_3 from this repo
+    pipe.tokenizer_3 = T5Tokenizer.from_pretrained(MODEL_ID, subfolder="tokenizer_3")
+
+    if DEVICE == "cuda":
+        pipe.to("cuda")
+        # Helps reduce VRAM spikes on smaller GPUs (slower but safer)
+        pipe.enable_model_cpu_offload()
+    else:
+        pipe.to("cpu")
+
+    PIPE = pipe
+    return PIPE
+
+
+def generate(mask_img, steps, cond_scale, seed):
+    if mask_img is None:
+        return None, "Please upload a semantic mask image."
+
+    pipe = load_pipe()
+
+    mask_pil = mask_img.convert("RGB")
+    orig_size = mask_pil.size
+
+    labels = mask_to_labels(mask_pil)
+    unique_ids = np.unique(labels).tolist()
+    organ_ids = [i for i in unique_ids if i != 0]
+
+    if organ_ids:
+        organ_names = [NAME_CLASS_DICT[i] for i in organ_ids]
+        prompt = "CT image containing " + " ".join(organ_names)
+    else:
+        organ_names = []
+        prompt = "CT image containing abdominal organs"
+
+    gen = torch.Generator(device=DEVICE).manual_seed(int(seed))
+
+    out = pipe(
+        prompt=prompt,
+        control_image=mask_pil,
+        height=128,
+        width=128,
+        num_inference_steps=int(steps),
+        generator=gen,
+        controlnet_conditioning_scale=float(cond_scale),
+    ).images[0]
+
+    out = out.resize(orig_size)
+    msg = f"Detected classes: {', '.join(['background'] + organ_names)}\nPrompt: {prompt}"
+    return out, msg
+
+
+with gr.Blocks(title="Mask2CT Atlas Demo (SD3 ControlNet)") as demo:
+    gr.Markdown(
+        """
+# Mask → CT (Atlas) demo
+Upload an **RGB semantic mask** (segmentation map) using the palette from the model card.
+The app auto-detects organs in the mask and generates a synthetic CT slice.
+"""
+    )
+
+    with gr.Row():
+        mask_in = gr.Image(label="Semantic mask (RGB)", type="pil")
+        ct_out = gr.Image(label="Generated CT", type="pil")
+
+    with gr.Row():
+        steps = gr.Slider(10, 80, value=50, step=1, label="Inference steps")
+        cond = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="ControlNet conditioning scale")
+        seed = gr.Number(value=1, precision=0, label="Seed")
+
+    btn = gr.Button("Generate")
+    info = gr.Textbox(label="Info", lines=3)
+
+    btn.click(fn=generate, inputs=[mask_in, steps, cond, seed], outputs=[ct_out, info])
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,9 @@
+torch
+diffusers>=0.36.0
+transformers
+accelerate
+safetensors
+sentencepiece
+numpy
+pillow
+gradio