Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,33 @@
+# Base model weights (downloaded at runtime from HF Hub)
+scripts/models/Qwen/
+
+# LoRA checkpoints — keep only epoch-2 for the demo
+scripts/models/qwen_image_edit_chexpert_lora/epoch-0.safetensors
+scripts/models/qwen_image_edit_chexpert_lora/epoch-1.safetensors
+scripts/models/qwen_image_edit_chexpert_lora/epoch-3.safetensors
+scripts/models/qwen_image_edit_chexpert_lora/epoch-4.safetensors
+
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+*.egg
+
+# Environment
+.env
+.venv/
+venv/
+.cache/
+
+# IDE
+.vscode/
+.idea/
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Misc
+*.log

README.md

@@ -1,14 +1,31 @@
 ---
 title: SynthCXR
-emoji: 🌍
-colorFrom: yellow
-colorTo: green
+emoji: 🫁
+colorFrom: blue
+colorTo: purple
 sdk: gradio
-sdk_version: 6.9.0
-python_version: '3.12'
+sdk_version: "6.9.0"
 app_file: app.py
+hardware: zero-a10g
 pinned: false
-license: mit
+tags:
+  - medical-imaging
+  - chest-x-ray
+  - diffusion
+  - lora
+short_description: Controllable chest X-ray generation with anatomical masks
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🫁 SynthCXR · Chest X-Ray Generator
+
+Interactively resize anatomical mask components (heart, left lung, right lung) with sliders and generate realistic chest X-rays using a Qwen-Image-Edit model with LoRA fine-tuning on CheXpert.
+
+> **Zero GPU** — This Space uses HuggingFace ZeroGPU for dynamic GPU allocation. A GPU is acquired only during image generation and released immediately after.
+
+## Features
+
+- **Mask Scaling Sliders** — Real-time preview of organ masks scaled from 0× to 2×
+- **Condition Picker** — Select from 13 CheXpert pathologies with severity modifiers
+- **Demographics** — Configure patient age, sex, and radiograph view (AP/PA)
+- **CXR Generation** — Generate 512×512 chest X-rays conditioned on the modified mask
+- **Progress Bar** — Real-time step-by-step progress during generation

app.py

@@ -0,0 +1,721 @@
+#!/usr/bin/env python3
+"""Gradio app for SynthCXR: interactive mask scaling and CXR generation."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import spaces
+
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+
+from synthcxr.constants import KNOWN_CONDITIONS
+from synthcxr.mask_utils import resolve_overlaps, scale_mask_channel
+from synthcxr.prompt import ConditionConfig, build_condition_prompt
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+BASE_DIR = Path(__file__).resolve().parent
+SAMPLE_MASKS_DIR = BASE_DIR / "static" / "sample_masks"
+LORA_DIR = BASE_DIR / "scripts" / "models" / "qwen_image_edit_chexpert_lora"
+
+# ---------------------------------------------------------------------------
+# Condition / severity choices
+# ---------------------------------------------------------------------------
+CONDITION_CHOICES = [
+    "enlarged_cardiomediastinum",
+    "cardiomegaly",
+    "atelectasis",
+    "pneumothorax",
+    "pleural_effusion",
+]
+SEVERITY_CHOICES = ["(none)", "mild", "moderate", "severe"]
+
+# ---------------------------------------------------------------------------
+# Pipeline (lazy-loaded once)
+# ---------------------------------------------------------------------------
+_pipe = None
+
+
+def get_pipeline():
+    """Load the diffusion pipeline + LoRA weights into GPU memory (once)."""
+    global _pipe
+    if _pipe is not None:
+        return _pipe
+
+    from synthcxr.pipeline import load_lora_weights, load_pipeline
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16
+
+    # VRAM_LIMIT (in GB): enables model offloading for memory-constrained GPUs
+    vram_limit_str = os.environ.get("VRAM_LIMIT", "")
+    vram_limit = float(vram_limit_str) if vram_limit_str else None
+
+    print(f"[INFO] Loading QwenImagePipeline (device={device}, dtype={dtype}, vram_limit={vram_limit}) …")
+    _pipe = load_pipeline(device, dtype, vram_limit=vram_limit)
+
+    # LORA_EPOCH env var: which epoch checkpoint to load (default: 2)
+    lora_epoch = os.environ.get("LORA_EPOCH", "2")
+    lora = LORA_DIR / f"epoch-{lora_epoch}.safetensors"
+
+    if not lora.exists():
+        # Try step-based checkpoints or any available .safetensors
+        candidates = sorted(LORA_DIR.glob("*.safetensors")) if LORA_DIR.exists() else []
+        if candidates:
+            lora = candidates[-1]
+            print(f"[WARN] epoch-{lora_epoch} not found, falling back to {lora.name}")
+        else:
+            print("[WARN] No LoRA checkpoint found – running base model only.")
+            return _pipe
+
+    print(f"[INFO] Loading LoRA from {lora}")
+    load_lora_weights(_pipe, lora)
+    print("[INFO] Pipeline ready.")
+    return _pipe
+
+
+# ---------------------------------------------------------------------------
+# Sample masks
+# ---------------------------------------------------------------------------
+def get_sample_masks() -> list[str]:
+    """Return paths of bundled sample masks."""
+    if not SAMPLE_MASKS_DIR.exists():
+        return []
+    return sorted(str(p) for p in SAMPLE_MASKS_DIR.glob("*.png"))
+
+
+# ---------------------------------------------------------------------------
+# Core functions
+# ---------------------------------------------------------------------------
+
+def apply_mask_scaling(
+    mask_array: np.ndarray,
+    heart_scale: float,
+    left_lung_scale: float,
+    right_lung_scale: float,
+) -> np.ndarray:
+    """Scale mask channels and resolve overlaps."""
+    if heart_scale != 1.0:
+        mask_array = scale_mask_channel(mask_array, channel=2, scale_factor=heart_scale)
+    if left_lung_scale != 1.0:
+        mask_array = scale_mask_channel(mask_array, channel=0, scale_factor=left_lung_scale)
+    if right_lung_scale != 1.0:
+        mask_array = scale_mask_channel(mask_array, channel=1, scale_factor=right_lung_scale)
+    return resolve_overlaps(mask_array, priority=(2, 0, 1))
+
+
+def preview_mask(
+    mask_image: np.ndarray | None,
+    heart_scale: float,
+    left_lung_scale: float,
+    right_lung_scale: float,
+) -> np.ndarray | None:
+    """Live mask preview callback."""
+    if mask_image is None:
+        return None
+    mask = np.array(Image.fromarray(mask_image).convert("RGB"))
+    scaled = apply_mask_scaling(mask, heart_scale, left_lung_scale, right_lung_scale)
+    return scaled
+
+
+def build_prompt_preview(
+    conditions: list[str],
+    severity: str,
+    age: int,
+    sex: str,
+    view: str,
+) -> str:
+    """Build the prompt text for preview."""
+    cond = ConditionConfig(
+        name="preview",
+        conditions=conditions or [],
+        age=age,
+        sex=sex,
+        view=view,
+        severity=severity if severity != "(none)" else None,
+    )
+    return build_condition_prompt(cond)
+
+
+@spaces.GPU(duration=120)
+def generate_cxr(
+    mask_image: np.ndarray | None,
+    heart_scale: float,
+    left_lung_scale: float,
+    right_lung_scale: float,
+    conditions: list[str],
+    severity: str,
+    age: int,
+    sex: str,
+    view: str,
+    num_steps: int,
+    cfg_scale: float,
+    seed: int,
+    preview_every: int = 10,
+    progress=gr.Progress(),
+):
+    """Generate a CXR, yielding intermediate previews every N steps."""
+    if mask_image is None:
+        raise gr.Error("Please select or upload a mask first.")
+
+    pipe = get_pipeline()
+    if pipe is None:
+        raise gr.Error("Pipeline not loaded. GPU may be unavailable.")
+
+    # Prepare mask
+    mask = np.array(Image.fromarray(mask_image).convert("RGB"))
+    scaled = apply_mask_scaling(mask, heart_scale, left_lung_scale, right_lung_scale)
+    edit_image = Image.fromarray(scaled)
+
+    # Build prompt
+    cond = ConditionConfig(
+        name="web_ui",
+        conditions=conditions or [],
+        age=age,
+        sex=sex,
+        view=view,
+        severity=severity if severity != "(none)" else None,
+    )
+    prompt = build_condition_prompt(cond)
+
+    # Intermediate preview collector
+    previews: list[Image.Image] = []
+
+    class StepCallback:
+        """Custom tqdm-like wrapper that decodes latents every N steps."""
+        def __init__(self, iterable):
+            self._iterable = iterable
+            self._step = 0
+
+        def __iter__(self):
+            for item in self._iterable:
+                progress(self._step / num_steps, desc="Generating CXR...")
+                yield item
+                self._step += 1
+                if (
+                    preview_every > 0
+                    and self._step % preview_every == 0
+                    and self._step < num_steps
+                    and "latents" in _shared_ref
+                ):
+                    try:
+                        with torch.no_grad():
+                            latents = _shared_ref["latents"]
+                            decoded = pipe.vae.decode(
+                                latents,
+                                device=pipe.device,
+                                tiled=False,
+                            )
+                            img = pipe.vae_output_to_image(decoded)
+                            previews.append(img)
+                    except Exception:
+                        pass  # skip preview on error
+
+        def __len__(self):
+            return len(self._iterable)
+
+    # We patch the pipeline's __call__ to capture inputs_shared reference.
+    # The pipeline stores latents in inputs_shared["latents"] during denoising.
+    _shared_ref: dict = {}
+    _orig_unit_runner = pipe.unit_runner.__class__.__call__
+
+    def _patched_runner(self_runner, unit, p, inputs_shared, inputs_posi, inputs_nega):
+        _shared_ref.update(inputs_shared)
+        return _orig_unit_runner(self_runner, unit, p, inputs_shared, inputs_posi, inputs_nega)
+
+    pipe.unit_runner.__class__.__call__ = _patched_runner
+
+    try:
+        image = pipe(
+            prompt=prompt,
+            edit_image=edit_image,
+            height=512,
+            width=512,
+            num_inference_steps=num_steps,
+            seed=seed,
+            rand_device=pipe.device,
+            cfg_scale=cfg_scale,
+            edit_image_auto_resize=True,
+            zero_cond_t=True,
+            progress_bar_cmd=StepCallback,
+        )
+    finally:
+        # Restore original runner
+        pipe.unit_runner.__class__.__call__ = _orig_unit_runner
+
+    # Yield all collected previews, then the final image
+    for preview in previews:
+        yield preview
+    yield image
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+CUSTOM_CSS = """
+/* ── Layout ── */
+.gradio-container {
+    max-width: 1280px !important;
+    margin: 0 auto !important;
+}
+
+/* ── Radial gradient background ── */
+.main { 
+    background: 
+        radial-gradient(ellipse 80% 50% at 10% 20%, rgba(99,102,241,0.07), transparent),
+        radial-gradient(ellipse 60% 40% at 85% 75%, rgba(59,130,246,0.05), transparent) !important;
+}
+
+/* ── Header ── */
+#component-0 h1 {
+    text-align: center;
+    font-size: 2.2rem !important;
+    font-weight: 800 !important;
+    letter-spacing: -0.5px;
+    background: linear-gradient(135deg, #818cf8, #60a5fa, #818cf8);
+    background-size: 200% 200%;
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    animation: gradientShift 4s ease-in-out infinite;
+    padding-bottom: 4px !important;
+}
+#component-0 p {
+    text-align: center;
+    color: #94a3b8 !important;
+    font-size: 0.95rem;
+}
+
+@keyframes gradientShift {
+    0%, 100% { background-position: 0% 50%; }
+    50% { background-position: 100% 50%; }
+}
+
+/* ── Glass panels ── */
+.block {
+    border: 1px solid rgba(99,115,146,0.15) !important;
+    border-radius: 16px !important;
+    backdrop-filter: blur(12px);
+    transition: border-color 0.3s ease, box-shadow 0.3s ease !important;
+}
+.block:hover {
+    border-color: rgba(99,102,241,0.25) !important;
+    box-shadow: 0 0 20px rgba(99,102,241,0.06) !important;
+}
+
+/* ── Section headings ── */
+.markdown h3 {
+    font-size: 0.78rem !important;
+    font-weight: 700 !important;
+    text-transform: uppercase;
+    letter-spacing: 1.2px;
+    color: #64748b !important;
+    border-bottom: 1px solid rgba(99,115,146,0.12);
+    padding-bottom: 8px !important;
+    margin-bottom: 12px !important;
+}
+
+/* ── Slider styling ── */
+input[type="range"] {
+    height: 6px !important;
+    border-radius: 3px !important;
+    background: #1e293b !important;
+}
+input[type="range"]::-webkit-slider-thumb {
+    width: 18px !important;
+    height: 18px !important;
+    border-radius: 50% !important;
+    border: 2.5px solid #0a0e17 !important;
+    transition: transform 0.2s ease, box-shadow 0.2s ease !important;
+}
+input[type="range"]::-webkit-slider-thumb:hover {
+    transform: scale(1.2) !important;
+}
+
+/* Slider labels */
+.block label span {
+    font-weight: 500 !important;
+    font-size: 0.88rem !important;
+}
+.block .rangeSlider_value {
+    font-variant-numeric: tabular-nums;
+    font-weight: 600 !important;
+}
+
+/* ── Image panels ── */
+.image-frame img, .image-container img {
+    border-radius: 10px !important;
+    transition: opacity 0.3s ease !important;
+}
+.image-container {
+    background: rgba(0,0,0,0.2) !important;
+    border-radius: 12px !important;
+    min-height: 380px;
+}
+
+/* ── Generate button ── */
+.primary {
+    background: linear-gradient(135deg, #6366f1, #4f46e5, #6366f1) !important;
+    background-size: 200% 200% !important;
+    border: none !important;
+    border-radius: 12px !important;
+    padding: 14px 24px !important;
+    font-weight: 700 !important;
+    font-size: 1rem !important;
+    letter-spacing: 0.3px;
+    transition: all 0.3s cubic-bezier(0.4,0,0.2,1) !important;
+    position: relative;
+    overflow: hidden;
+}
+.primary:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 8px 25px rgba(99,102,241,0.4) !important;
+    animation: btnShimmer 1.5s ease-in-out infinite !important;
+}
+.primary:active {
+    transform: translateY(0) !important;
+}
+@keyframes btnShimmer {
+    0%, 100% { background-position: 0% 50%; }
+    50% { background-position: 100% 50%; }
+}
+
+/* ── Secondary buttons ── */
+.secondary {
+    border: 1px solid rgba(99,115,146,0.2) !important;
+    border-radius: 10px !important;
+    background: transparent !important;
+    color: #94a3b8 !important;
+    transition: all 0.25s ease !important;
+}
+.secondary:hover {
+    border-color: rgba(99,102,241,0.4) !important;
+    color: #e2e8f0 !important;
+    background: rgba(99,102,241,0.06) !important;
+}
+
+/* ── Prompt preview ── */
+textarea[readonly], .prose {
+    font-family: 'JetBrains Mono', 'Fira Code', monospace !important;
+    font-size: 0.8rem !important;
+    line-height: 1.6 !important;
+    color: #64748b !important;
+    background: rgba(0,0,0,0.25) !important;
+    border-radius: 10px !important;
+}
+
+/* ── Checkboxes ── */
+.checkbox-group label {
+    border-radius: 20px !important;
+    padding: 4px 12px !important;
+    font-size: 0.8rem !important;
+    transition: all 0.2s ease !important;
+    border: 1px solid rgba(99,115,146,0.15) !important;
+    color: #e2e8f0 !important;
+    background: rgba(17,24,39,0.75) !important;
+}
+.checkbox-group label span {
+    color: #e2e8f0 !important;
+}
+.checkbox-group label:hover {
+    border-color: rgba(99,102,241,0.35) !important;
+    background: rgba(30,41,59,0.9) !important;
+}
+.checkbox-group input:checked + label,
+.checkbox-group label.selected {
+    background: rgba(99,102,241,0.15) !important;
+    border-color: rgba(99,102,241,0.4) !important;
+    color: #c7d2fe !important;
+}
+
+/* ── Dropdowns & inputs ── */
+select, input[type="number"] {
+    border-radius: 10px !important;
+    border: 1px solid rgba(99,115,146,0.15) !important;
+    transition: border-color 0.25s ease !important;
+    font-size: 0.88rem !important;
+}
+select:focus, input[type="number"]:focus {
+    border-color: rgba(99,102,241,0.5) !important;
+    box-shadow: 0 0 0 2px rgba(99,102,241,0.1) !important;
+}
+
+/* ── Accordion ── */
+.accordion {
+    border: 1px solid rgba(99,115,146,0.1) !important;
+    border-radius: 12px !important;
+    background: rgba(0,0,0,0.15) !important;
+}
+.accordion > .label-wrap {
+    font-size: 0.82rem !important;
+    color: #64748b !important;
+    font-weight: 500 !important;
+}
+
+/* ── Examples gallery ── */
+.gallery-item {
+    border-radius: 10px !important;
+    border: 2px solid rgba(99,115,146,0.15) !important;
+    transition: all 0.25s ease !important;
+    overflow: hidden;
+}
+.gallery-item:hover {
+    border-color: rgba(99,102,241,0.4) !important;
+    transform: scale(1.04);
+    box-shadow: 0 4px 16px rgba(99,102,241,0.15) !important;
+}
+
+/* ── Scrollbar ── */
+::-webkit-scrollbar { width: 6px; }
+::-webkit-scrollbar-track { background: transparent; }
+::-webkit-scrollbar-thumb {
+    background: rgba(99,115,146,0.25);
+    border-radius: 3px;
+}
+::-webkit-scrollbar-thumb:hover { background: rgba(99,115,146,0.4); }
+
+/* ── Footer spacing ── */
+.gradio-container > .main > .wrap:last-child { padding-bottom: 40px !important; }
+"""
+
+sample_paths = get_sample_masks()
+
+THEME = gr.themes.Base(
+    primary_hue=gr.themes.colors.indigo,
+    secondary_hue=gr.themes.colors.slate,
+    neutral_hue=gr.themes.colors.slate,
+    font=gr.themes.GoogleFont("Inter"),
+    font_mono=gr.themes.GoogleFont("JetBrains Mono"),
+    radius_size=gr.themes.sizes.radius_lg,
+    spacing_size=gr.themes.sizes.spacing_md,
+).set(
+    # Background
+    body_background_fill="#0a0e17",
+    body_background_fill_dark="#0a0e17",
+    # Panels
+    block_background_fill="rgba(17,24,39,0.75)",
+    block_background_fill_dark="rgba(17,24,39,0.75)",
+    block_border_color="rgba(99,115,146,0.15)",
+    block_border_color_dark="rgba(99,115,146,0.15)",
+    block_shadow="0 4px 24px rgba(0,0,0,0.2)",
+    block_shadow_dark="0 4px 24px rgba(0,0,0,0.2)",
+    # Inputs
+    input_background_fill="#131b2e",
+    input_background_fill_dark="#131b2e",
+    input_border_color="rgba(99,115,146,0.15)",
+    input_border_color_dark="rgba(99,115,146,0.15)",
+    # Buttons
+    button_primary_background_fill="linear-gradient(135deg, #6366f1, #4f46e5)",
+    button_primary_background_fill_dark="linear-gradient(135deg, #6366f1, #4f46e5)",
+    button_primary_text_color="white",
+    button_primary_text_color_dark="white",
+    button_primary_shadow="0 4px 14px rgba(99,102,241,0.25)",
+    button_primary_shadow_dark="0 4px 14px rgba(99,102,241,0.25)",
+    # Text
+    body_text_color="#e2e8f0",
+    body_text_color_dark="#e2e8f0",
+    body_text_color_subdued="#94a3b8",
+    body_text_color_subdued_dark="#94a3b8",
+    # Labels
+    block_label_text_color="#94a3b8",
+    block_label_text_color_dark="#94a3b8",
+    block_title_text_color="#cbd5e1",
+    block_title_text_color_dark="#cbd5e1",
+    # Borders
+    border_color_primary="rgba(99,102,241,0.4)",
+    border_color_primary_dark="rgba(99,102,241,0.4)",
+)
+
+with gr.Blocks(
+    title="SynthCXR · Chest X-Ray Generator",
+) as demo:
+
+    gr.Markdown(
+        "# 🫁 SynthCXR\n"
+        "Interactively resize anatomical masks and generate realistic chest X-rays"
+    )
+
+    with gr.Row():
+
+        # ── Left column: Controls ──
+        with gr.Column(scale=1):
+
+            # Mask input
+            gr.Markdown("### Select Mask")
+            mask_input = gr.Image(
+                label="Conditioning Mask",
+                type="numpy",
+                sources=["upload"],
+                height=240,
+            )
+
+            # Sample mask gallery
+            if sample_paths:
+                sample_gallery = gr.Examples(
+                    examples=sample_paths,
+                    inputs=mask_input,
+                    label="Sample Masks",
+                )
+
+            # Sliders
+            gr.Markdown("### Mask Scaling")
+            heart_slider = gr.Slider(
+                minimum=0.0, maximum=2.0, step=0.05, value=1.0,
+                label="💙 Heart Scale",
+            )
+            left_lung_slider = gr.Slider(
+                minimum=0.0, maximum=2.0, step=0.05, value=1.0,
+                label="🔴 Left Lung Scale",
+            )
+            right_lung_slider = gr.Slider(
+                minimum=0.0, maximum=2.0, step=0.05, value=1.0,
+                label="🟢 Right Lung Scale",
+            )
+            reset_btn = gr.Button("↺ Reset Scales", variant="secondary", size="sm")
+
+            # Conditions
+            gr.Markdown("### Conditions")
+            conditions_select = gr.CheckboxGroup(
+                choices=CONDITION_CHOICES,
+                label="Pathologies",
+            )
+            with gr.Row():
+                severity_select = gr.Radio(
+                    choices=SEVERITY_CHOICES, value="(none)", label="Severity",
+                )
+                view_select = gr.Radio(
+                    choices=["AP", "PA"], value="AP", label="View",
+                )
+            with gr.Row():
+                age_input = gr.Number(value=45, label="Age", minimum=0, maximum=120, precision=0)
+                sex_select = gr.Radio(
+                    choices=["male", "female"], value="male", label="Sex",
+                )
+
+            # Advanced
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    steps_input = gr.Number(value=50, label="Steps", minimum=1, maximum=100, precision=0)
+                    cfg_input = gr.Number(value=4.0, label="CFG Scale", minimum=1.0, maximum=20.0)
+                with gr.Row():
+                    seed_input = gr.Number(value=42, label="Seed", minimum=0, precision=0)
+                    preview_every_input = gr.Number(value=10, label="Preview Every N Steps", minimum=0, maximum=50, precision=0)
+
+        # ── Right column: Outputs ──
+        with gr.Column(scale=2):
+
+            with gr.Row():
+                mask_preview = gr.Image(
+                    label="Scaled Mask Preview",
+                    type="numpy",
+                    interactive=False,
+                    height=400,
+                )
+                cxr_output = gr.Image(
+                    label="Generated Chest X-Ray",
+                    type="pil",
+                    interactive=False,
+                    height=400,
+                )
+
+            # Prompt preview
+            prompt_preview = gr.Textbox(
+                label="Prompt Preview",
+                interactive=False,
+                lines=3,
+            )
+
+            generate_btn = gr.Button("⚡ Generate CXR", variant="primary", size="lg")
+
+    # ── Event wiring ──
+
+    # Live mask preview on any slider / mask change
+    slider_inputs = [mask_input, heart_slider, left_lung_slider, right_lung_slider]
+
+    mask_input.change(preview_mask, inputs=slider_inputs, outputs=mask_preview)
+    heart_slider.change(preview_mask, inputs=slider_inputs, outputs=mask_preview)
+    left_lung_slider.change(preview_mask, inputs=slider_inputs, outputs=mask_preview)
+    right_lung_slider.change(preview_mask, inputs=slider_inputs, outputs=mask_preview)
+
+    # Reset sliders
+    def reset_scales():
+        return 1.0, 1.0, 1.0
+
+    reset_btn.click(
+        reset_scales,
+        outputs=[heart_slider, left_lung_slider, right_lung_slider],
+    )
+
+    # Auto-adjust sliders when conditions change
+    _CONDITION_SCALE_MAP = {
+        # condition_key: (heart_delta, lung_delta)
+        "cardiomegaly":                (+0.35, 0.0),
+        "enlarged_cardiomediastinum":  (+0.25, 0.0),
+        "atelectasis":                 (0.0, -0.25),
+        "pneumothorax":                (0.0, -0.30),
+        "pleural_effusion":            (0.0, -0.20),
+    }
+    _SEVERITY_MULTIPLIER = {
+        "(none)": 1.0,
+        "mild": 0.6,
+        "moderate": 1.0,
+        "severe": 1.5,
+    }
+
+    def sync_sliders(conditions: list[str], severity: str):
+        """Set slider values based on selected conditions + severity."""
+        heart = 1.0
+        lung = 1.0
+        mult = _SEVERITY_MULTIPLIER.get(severity, 1.0)
+        for cond in (conditions or []):
+            h_delta, l_delta = _CONDITION_SCALE_MAP.get(cond, (0.0, 0.0))
+            heart += h_delta * mult
+            lung += l_delta * mult
+        # Clamp to slider range [0.0, 2.0]
+        heart = round(max(0.0, min(2.0, heart)), 2)
+        lung = round(max(0.0, min(2.0, lung)), 2)
+        return heart, lung, lung
+
+    conditions_select.change(
+        sync_sliders,
+        inputs=[conditions_select, severity_select],
+        outputs=[heart_slider, left_lung_slider, right_lung_slider],
+    )
+    severity_select.change(
+        sync_sliders,
+        inputs=[conditions_select, severity_select],
+        outputs=[heart_slider, left_lung_slider, right_lung_slider],
+    )
+
+    # Prompt preview on config change
+    prompt_inputs = [conditions_select, severity_select, age_input, sex_select, view_select]
+
+    for inp in prompt_inputs:
+        inp.change(build_prompt_preview, inputs=prompt_inputs, outputs=prompt_preview)
+
+    # Generate
+    generate_btn.click(
+        generate_cxr,
+        inputs=[
+            mask_input,
+            heart_slider, left_lung_slider, right_lung_slider,
+            conditions_select, severity_select,
+            age_input, sex_select, view_select,
+            steps_input, cfg_input, seed_input,
+            preview_every_input,
+        ],
+        outputs=cxr_output,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Launch (module-level for HuggingFace Spaces compatibility)
+# ---------------------------------------------------------------------------
+demo.launch(theme=THEME, css=CUSTOM_CSS)

pyproject.toml

@@ -0,0 +1,20 @@
+[project]
+name = "synthcxr"
+version = "0.1.0"
+description = "Chest X-ray generation via Qwen-Image-Edit LoRA fine-tuning"
+requires-python = ">=3.10.1"
+dependencies = [
+    "diffsynth>=2.0.4",
+    "fastapi[standard]>=0.135.1",
+    "gradio>=6.8.0",
+    "python-multipart>=0.0.22",
+    "scipy",
+    "uvicorn[standard]>=0.41.0",
+]
+
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]

requirements.txt

@@ -0,0 +1,7 @@
+gradio>=4.0
+diffsynth>=2.0.4
+spaces
+scipy
+Pillow
+numpy
+torch

scripts/models/qwen_image_edit_chexpert_lora/epoch-2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fef90b53ae95c9628efe14b0919f7be7e291ec9f80677a3f2ed509ebccca1c05
+size 472047184

scripts/models/qwen_image_edit_chexpert_lora/latest_checkpoint.json

@@ -0,0 +1 @@
+{"path": "./models/qwen_image_edit_chexpert_lora/checkpoint-step233240", "epoch_id": 4, "global_step": 233240}

synthcxr/__init__.py

@@ -0,0 +1 @@
+"""SynthCXR: Chest X-ray generation via Qwen-Image-Edit LoRA fine-tuning."""

synthcxr/constants.py

@@ -0,0 +1,54 @@
+"""Shared constants for SynthCXR: disease labels, condition maps, severity modifiers."""
+
+from __future__ import annotations
+
+# CheXpert label column names -> natural-language descriptions used in prompts.
+# Used by both dataset preparation and inference scripts.
+LABEL_TEXT: dict[str, str] = {
+    "Enlarged Cardiomediastinum": "enlarged cardiomediastinum",
+    "Cardiomegaly": "cardiomegaly",
+    "Lung Opacity": "diffuse lung opacity",
+    "Lung Lesion": "discrete lung lesion",
+    "Edema": "pulmonary edema",
+    "Consolidation": "parenchymal consolidation",
+    "Pneumonia": "findings compatible with pneumonia",
+    "Atelectasis": "atelectasis",
+    "Pneumothorax": "pneumothorax",
+    "Pleural Effusion": "pleural effusion",
+    "Pleural Other": "other pleural abnormality",
+    "Fracture": "possible fracture",
+    "Support Devices": "support devices in place",
+}
+
+# Snake_case keys for config files -> natural-language descriptions.
+KNOWN_CONDITIONS: dict[str, str] = {
+    "enlarged_cardiomediastinum": "enlarged cardiomediastinum",
+    "cardiomegaly": "cardiomegaly",
+    "lung_opacity": "diffuse lung opacity",
+    "lung_lesion": "discrete lung lesion",
+    "edema": "pulmonary edema",
+    "consolidation": "parenchymal consolidation",
+    "pneumonia": "findings compatible with pneumonia",
+    "atelectasis": "atelectasis",
+    "pneumothorax": "pneumothorax",
+    "pleural_effusion": "pleural effusion",
+    "pleural_other": "other pleural abnormality",
+    "fracture": "possible fracture",
+    "support_devices": "support devices in place",
+}
+
+SEVERITY_MODIFIERS: dict[str, str] = {
+    "mild": "mild",
+    "moderate": "moderate",
+    "severe": "severe",
+    "small": "small",
+    "large": "large",
+    "very_small": "very small",
+    "very_large": "very large",
+    "minimal": "minimal",
+    "significant": "significant",
+}
+
+DEFAULT_MODEL_ID = "Qwen/Qwen-Image-Edit-2511"
+TEXT_ENCODER_MODEL_ID = "Qwen/Qwen-Image"
+PROCESSOR_MODEL_ID = "Qwen/Qwen-Image-Edit"

synthcxr/mask_utils.py

@@ -0,0 +1,93 @@
+"""Mask manipulation: scaling organ regions and resolving overlaps."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+from scipy import ndimage
+from scipy.ndimage import map_coordinates
+
+
+def resolve_overlaps(
+    mask: np.ndarray,
+    priority: tuple[int, int, int] = (2, 0, 1),
+    threshold: int = 10,
+) -> np.ndarray:
+    """Assign overlapping pixels to the highest-priority channel.
+
+    Default priority: heart (2) > left lung (0) > right lung (1).
+    """
+    result = mask.copy()
+    active = mask > threshold
+    overlap_mask = active.sum(axis=2) > 1
+    if not overlap_mask.any():
+        return result
+
+    for y, x in zip(*np.where(overlap_mask)):
+        active_channels = [ch for ch in range(3) if mask[y, x, ch] > threshold]
+        best = min(active_channels, key=lambda ch: priority.index(ch))
+        for ch in active_channels:
+            if ch != best:
+                result[y, x, ch] = 0
+    return result
+
+
+def scale_mask_channel(
+    mask: np.ndarray,
+    channel: int,
+    scale_factor: float,
+    threshold: int = 10,
+) -> np.ndarray:
+    """Scale a single channel's region around its centroid.
+
+    ``channel``: 0 = left lung (red), 1 = right lung (green), 2 = heart (blue).
+    """
+    result = mask.copy()
+    channel_data = mask[:, :, channel]
+    binary = channel_data > threshold
+    if not binary.any():
+        return result
+
+    cy, cx = ndimage.center_of_mass(binary)
+    h, w = mask.shape[:2]
+    y_coords, x_coords = np.mgrid[0:h, 0:w]
+    y_t = ((y_coords - cy) / scale_factor + cy).astype(np.float32)
+    x_t = ((x_coords - cx) / scale_factor + cx).astype(np.float32)
+
+    result[:, :, channel] = 0
+    scaled = map_coordinates(
+        channel_data.astype(np.float32),
+        [y_t, x_t],
+        order=1,
+        mode="constant",
+        cval=0,
+    )
+    result[:, :, channel] = np.clip(scaled, 0, 255).astype(np.uint8)
+    return result
+
+
+def modify_mask(
+    input_path: Path,
+    output_path: Path,
+    heart_scale: float = 1.0,
+    left_lung_scale: float = 1.0,
+    right_lung_scale: float = 1.0,
+) -> None:
+    """Load a conditioning mask, apply scale factors, and save."""
+    with Image.open(input_path) as img:
+        mask = np.array(img.convert("RGB"))
+
+    if left_lung_scale != 1.0:
+        mask = scale_mask_channel(mask, channel=0, scale_factor=left_lung_scale)
+    if right_lung_scale != 1.0:
+        mask = scale_mask_channel(mask, channel=1, scale_factor=right_lung_scale)
+    if heart_scale != 1.0:
+        mask = scale_mask_channel(mask, channel=2, scale_factor=heart_scale)
+
+    mask = resolve_overlaps(mask, priority=(2, 0, 1))
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    Image.fromarray(mask).save(output_path)
+    print(f"[INFO] Saved modified mask to {output_path}")

synthcxr/pipeline.py

@@ -0,0 +1,120 @@
+"""Pipeline loading, LoRA weight management, and image I/O helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Sequence
+
+import numpy as np
+import torch
+from PIL import Image
+
+from diffsynth.core import ModelConfig
+from diffsynth.pipelines.qwen_image import QwenImagePipeline
+
+from .constants import DEFAULT_MODEL_ID, PROCESSOR_MODEL_ID, TEXT_ENCODER_MODEL_ID
+from .mask_utils import resolve_overlaps, scale_mask_channel
+
+
+@dataclass
+class SampleSpec:
+    """A single validation/inference sample."""
+
+    prompt: str
+    mask_paths: list[Path]
+    identifier: str
+    image_path: Path | None
+    original_prompt: str = ""
+
+
+def load_pipeline(
+    device: str,
+    torch_dtype: torch.dtype,
+    model_id: str = DEFAULT_MODEL_ID,
+    vram_limit: float | None = None,
+) -> QwenImagePipeline:
+    """Instantiate a ``QwenImagePipeline``, downloading weights from HF Hub."""
+    model_configs = [
+        ModelConfig(
+            model_id=model_id,
+            origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors",
+        ),
+        ModelConfig(
+            model_id=TEXT_ENCODER_MODEL_ID,
+            origin_file_pattern="text_encoder/model*.safetensors",
+        ),
+        ModelConfig(
+            model_id=TEXT_ENCODER_MODEL_ID,
+            origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
+        ),
+    ]
+    tokenizer_config = ModelConfig(
+        model_id=TEXT_ENCODER_MODEL_ID,
+        origin_file_pattern="tokenizer/",
+    )
+    processor_config = ModelConfig(
+        model_id=PROCESSOR_MODEL_ID,
+        origin_file_pattern="processor/",
+    )
+    pipe = QwenImagePipeline.from_pretrained(
+        torch_dtype=torch_dtype,
+        device=device,
+        model_configs=model_configs,
+        tokenizer_config=tokenizer_config,
+        processor_config=processor_config,
+        vram_limit=vram_limit,
+    )
+    return pipe
+
+
+def load_lora_weights(pipe: QwenImagePipeline, checkpoint: Path) -> None:
+    """Load a LoRA checkpoint into an existing pipeline."""
+    pipe.clear_lora()
+    pipe.load_lora(pipe.dit, lora_config=str(checkpoint))
+
+
+def load_edit_images(
+    paths: Sequence[Path],
+    *,
+    heart_scale: float = 1.0,
+    left_lung_scale: float = 1.0,
+    right_lung_scale: float = 1.0,
+) -> Image.Image | list[Image.Image]:
+    """Load conditioning mask image(s), optionally rescaling organ regions."""
+    images: list[Image.Image] = []
+    needs_modification = (
+        heart_scale != 1.0 or left_lung_scale != 1.0 or right_lung_scale != 1.0
+    )
+    for path in paths:
+        with Image.open(path) as img:
+            if needs_modification:
+                mask = np.array(img.convert("RGB"))
+                if heart_scale != 1.0:
+                    mask = scale_mask_channel(mask, channel=2, scale_factor=heart_scale)
+                if left_lung_scale != 1.0:
+                    mask = scale_mask_channel(mask, channel=0, scale_factor=left_lung_scale)
+                if right_lung_scale != 1.0:
+                    mask = scale_mask_channel(mask, channel=1, scale_factor=right_lung_scale)
+                mask = resolve_overlaps(mask, priority=(2, 0, 1))
+                images.append(Image.fromarray(mask))
+            else:
+                images.append(img.convert("RGB"))
+    return images[0] if len(images) == 1 else images
+
+
+def export_original_images(
+    samples: Sequence[SampleSpec], output_dir: Path
+) -> None:
+    """Copy original CXR images into *output_dir*/original/ for comparison."""
+    original_dir = output_dir / "original"
+    original_dir.mkdir(parents=True, exist_ok=True)
+    for sample in samples:
+        if sample.image_path is None:
+            continue
+        dest = original_dir / f"{sample.identifier}.png"
+        if dest.exists():
+            continue
+        with Image.open(sample.image_path) as img:
+            img.convert("RGB").save(dest)
+        print(f"[INFO] Saved original to {dest}")

synthcxr/prompt.py

@@ -0,0 +1,80 @@
+"""Prompt builders for conditional inference."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from .constants import KNOWN_CONDITIONS, SEVERITY_MODIFIERS
+
+
+@dataclass
+class ConditionConfig:
+    """Configuration for a single inference run with specific conditions."""
+
+    name: str
+    conditions: list[str] = field(default_factory=list)
+    age: int | None = None
+    sex: str | None = None
+    view: str = "AP"
+    custom_prompt: str | None = None
+    severity: str | None = None
+    heart_scale: float = 1.0
+    left_lung_scale: float = 1.0
+    right_lung_scale: float = 1.0
+
+
+@dataclass
+class InferenceConfig:
+    """Top-level configuration for the condition-inference script."""
+
+    num_samples: int = 10
+    num_steps: int = 50
+    height: int = 512
+    width: int = 512
+    cfg_scale: float = 4.0
+    seed: int = 0
+    conditions: list[ConditionConfig] = field(default_factory=list)
+
+
+def build_condition_prompt(condition: ConditionConfig) -> str:
+    """Build a CheXpert-style prompt from a ``ConditionConfig``."""
+    if condition.custom_prompt:
+        return condition.custom_prompt
+
+    view = condition.view.upper() if condition.view else "AP"
+    age_str = f"{condition.age}-year-old" if condition.age else ""
+    sex_str = condition.sex.lower() if condition.sex else ""
+
+    if age_str and sex_str:
+        demographics = f"a {age_str} {sex_str} patient"
+    elif age_str:
+        demographics = f"a {age_str} patient"
+    elif sex_str:
+        demographics = f"a {sex_str} patient"
+    else:
+        demographics = "a patient"
+
+    pathologies: list[str] = []
+    severity_prefix = ""
+    if condition.severity and condition.severity in SEVERITY_MODIFIERS:
+        severity_prefix = SEVERITY_MODIFIERS[condition.severity] + " "
+
+    for cond_key in condition.conditions:
+        cond_text = KNOWN_CONDITIONS.get(cond_key.lower(), cond_key)
+        if severity_prefix and not pathologies:
+            pathologies.append(severity_prefix + cond_text)
+            severity_prefix = ""
+        else:
+            pathologies.append(cond_text)
+
+    pathology_str = (
+        f"with {', '.join(pathologies)}" if pathologies else "with no significant abnormality"
+    )
+
+    return (
+        f"frontal {view} chest radiograph of {demographics} {pathology_str}. "
+        "The conditioning mask image provides three channels "
+        "(red=left lung, green=right lung, blue=heart). "
+        "Reconstruct a CheXpert-style chest X-ray that aligns "
+        "with the segmentation and follows the described pathology."
+    )

synthcxr/utils.py

@@ -0,0 +1,27 @@
+"""Small shared helpers used across scripts."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Sequence
+
+
+def resolve_path(base: Path, maybe_relative: str) -> Path:
+    """Return *maybe_relative* as an absolute path, resolved against *base*."""
+    path = Path(maybe_relative).expanduser()
+    if path.is_absolute():
+        return path
+    return (base / path).resolve()
+
+
+def build_identifier(
+    record_path: str | None,
+    fallback_paths: Sequence[str],
+    sample_idx: int,
+) -> str:
+    """Build a filesystem-safe identifier from a metadata record."""
+    source = record_path or (fallback_paths[0] if fallback_paths else f"sample_{sample_idx}")
+    candidate = Path(source)
+    tail_parts = [part.replace(".", "-") for part in candidate.parts[-4:]]
+    slug = "_".join(tail_parts) if tail_parts else candidate.stem
+    return f"{sample_idx:03d}_{slug}"