Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """Gradio app for SynthCXR: interactive mask scaling and CXR generation.""" | |
| from __future__ import annotations | |
| import os | |
| os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" | |
| os.environ["DIFFSYNTH_DOWNLOAD_SOURCE"] = "huggingface" | |
| 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 loading (fresh on each @spaces.GPU call; model files cached on disk) | |
| # --------------------------------------------------------------------------- | |
| def load_fresh_pipeline(): | |
| """Load the pipeline + LoRA onto the *currently allocated* GPU. | |
| ZeroGPU deallocates GPU memory after each ``@spaces.GPU`` call, so we | |
| cannot cache tensors between calls. However, diffsynth caches the | |
| model files on disk (HF Hub cache), so only tensor loading happens | |
| here — not a full download. | |
| """ | |
| 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(): | |
| 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) | |
| 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, | |
| 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 = load_fresh_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 | |
| 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=40, label="Steps", minimum=1, maximum=100, precision=0) | |
| cfg_input = gr.Number(value=8.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) | |
| # ── 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, | |
| ], | |
| outputs=cxr_output, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Launch (module-level for HuggingFace Spaces compatibility) | |
| # --------------------------------------------------------------------------- | |
| demo.launch(theme=THEME, css=CUSTOM_CSS) | |