Polish diffusion UI and churning controls

app.py

@@ -0,0 +1,634 @@
+import os
+from pathlib import Path
+from typing import List, Optional
+
+os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
+
+import jax
+import numpy as np
+import gradio as gr
+import matplotlib.pyplot as plt
+
+from pipeline import INTEGRATORS, load_pipeline_assets, resolve_integrator, sample_batch
+
+N_SAMPLES = 5
+MAX_STEPS = 80
+DEFAULT_STEPS = 20
+ROOT_DIR = Path(__file__).parent
+LOGO_PATH = ROOT_DIR / "logo.png"
+LOGO_VALUE = str(LOGO_PATH) if LOGO_PATH.exists() else None
+DEFAULT_CHURN_RATE = 0.0
+DEFAULT_CHURN_MIN = 0.0
+DEFAULT_CHURN_MAX = 0.0
+DEFAULT_NOISE_INFLATION = 1.0
+MAX_NOISE_INFLATION = 1.02
+SUMMARY_PLACEHOLDER_HTML = """
+<div class="summary-card is-empty">
+  <div class="summary-title">Ready to sample</div>
+  <p>Select an integrator, adjust the controls, then generate digits to inspect their trajectories.</p>
+</div>
+""".strip()
+
+CUSTOM_CSS = """
+body {background: radial-gradient(circle at top left, #ffe8d5, #fff7f0 55%, #fdf1f8);}
+#hero {
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  gap: 1.5rem;
+  background: rgba(255, 255, 255, 0.85);
+  padding: 1.5rem 2rem;
+  border-radius: 18px;
+  box-shadow: 0 18px 35px rgba(255, 135, 0, 0.15);
+  border: 1px solid rgba(255, 145, 0, 0.35);
+}
+.hero-logo img {max-width: 320px; width: 100%; object-fit: contain;}
+.hero-copy {font-size: 1.05rem !important; color: #7a3b09;}
+.control-card {
+  background: rgba(255, 255, 255, 0.92);
+  border-radius: 16px;
+  padding: 1.25rem;
+  border: 1px solid rgba(255, 166, 77, 0.35);
+  box-shadow: 0 14px 30px rgba(255, 140, 0, 0.12);
+}
+.generate-button button {
+  background: linear-gradient(135deg, #ff7e00, #ffb347);
+  color: #fff;
+  font-weight: 600;
+  border-radius: 12px;
+  box-shadow: 0 10px 20px rgba(255, 126, 0, 0.25);
+}
+.generate-button button:hover {filter: brightness(1.05);}
+.control-heading {
+  font-weight: 600;
+  color: #7a3b09;
+  margin-bottom: 0.6rem !important;
+}
+.plot-card {
+  background: rgba(255, 255, 255, 0.88);
+  border-radius: 16px;
+  padding: 1rem;
+  border: 1px solid rgba(255, 166, 77, 0.35);
+  box-shadow: inset 0 0 0 1px rgba(255, 255, 255, 0.4), 0 12px 28px rgba(255, 145, 0, 0.18);
+}
+.details-card {
+  border: none;
+  padding: 0;
+}
+.summary-card {
+  background: rgba(255, 255, 255, 0.9);
+  border-radius: 14px;
+  padding: 1.1rem 1.25rem;
+  border: 1px solid rgba(255, 166, 77, 0.35);
+  box-shadow: 0 12px 26px rgba(255, 145, 0, 0.16);
+  display: grid;
+  gap: 0.85rem;
+}
+.summary-card.is-empty {
+  border-style: dashed;
+  box-shadow: none;
+}
+.summary-title {
+  font-weight: 600;
+  font-size: 1.05rem;
+  color: #7a3b09;
+}
+.summary-section {
+  display: grid;
+  gap: 0.45rem;
+}
+.summary-grid {
+  display: grid;
+  grid-template-columns: repeat(auto-fit, minmax(120px, 1fr));
+  gap: 0.4rem;
+}
+.summary-pill {
+  background: rgba(255, 245, 233, 0.95);
+  border: 1px solid rgba(255, 166, 77, 0.45);
+  border-radius: 999px;
+  padding: 0.35rem 0.75rem;
+  font-size: 0.85rem;
+  display: inline-flex;
+  align-items: center;
+  gap: 0.35rem;
+  color: #7a3b09;
+  justify-content: center;
+}
+.summary-pill strong {font-weight: 600;}
+.summary-pill.integrator {
+  background: rgba(255, 231, 206, 0.95);
+  border-color: rgba(255, 160, 72, 0.65);
+  font-weight: 600;
+}
+.summary-divider {
+  border: none;
+  border-top: 1px dashed rgba(255, 166, 77, 0.4);
+  margin: 0.2rem 0;
+}
+.accordion-card {
+  --tw-border-opacity: 0.45;
+  border: 1px dashed rgba(255, 166, 77, 0.45) !important;
+  border-radius: 14px !important;
+  background: rgba(255, 255, 255, 0.88) !important;
+}
+.accordion-card > div:nth-child(1) {
+  font-weight: 600;
+  color: #7a3b09;
+}
+.churn-card {
+  margin-top: 0.75rem;
+  background: rgba(255, 255, 255, 0.85);
+  border-radius: 14px;
+  padding: 0.9rem 1rem 1.1rem;
+  border: 1px dashed rgba(255, 166, 77, 0.5);
+  box-shadow: inset 0 0 0 1px rgba(255, 255, 255, 0.55);
+}
+.churn-title {
+  font-size: 0.92rem !important;
+  color: #8a450f;
+  margin-bottom: 0.55rem !important;
+}
+.gallery-card {
+  background: rgba(255, 255, 255, 0.9);
+  border-radius: 16px;
+  padding: 0.3rem 0.4rem;
+  border: 1px solid rgba(255, 166, 77, 0.28);
+  box-shadow: inset 0 0 0 1px rgba(255, 255, 255, 0.25), 0 8px 18px rgba(255, 145, 0, 0.12);
+}
+.gallery-card [data-testid="upload-zone"] {
+  display: none !important;
+}
+.gallery-card .grid {
+  min-height: 180px;
+}
+.gallery-card img {
+  border-radius: 10px;
+  transition: transform 0.15s ease, box-shadow 0.15s ease;
+}
+.gallery-card img:hover {
+  transform: translateY(-2px);
+  box-shadow: 0 8px 14px rgba(255, 145, 0, 0.18);
+}
+.history-card {
+  background: rgba(255, 255, 255, 0.88);
+  border-radius: 16px;
+  padding: 0.9rem;
+  border: 1px solid rgba(255, 166, 77, 0.35);
+  box-shadow: inset 0 0 0 1px rgba(255, 255, 255, 0.35), 0 10px 22px rgba(255, 145, 0, 0.15);
+}
+.plot-title {
+  color: #7a3b09 !important;
+  text-align: center;
+  font-weight: 600 !important;
+  margin-bottom: 0.45rem !important;
+}
+.history-placeholder {
+  text-align: center;
+  color: #8a450f;
+  font-size: 0.9rem;
+  margin-top: 0.5rem;
+}
+.value-chip {
+  background: rgba(255, 231, 206, 0.9);
+  border-radius: 999px;
+  padding: 0.1rem 0.55rem;
+  font-size: 0.82rem;
+  margin-left: 0.4rem;
+  color: #84400e;
+}
+@media (max-width: 768px) {
+  .summary-grid {grid-template-columns: repeat(auto-fit, minmax(110px, 1fr));}
+  .gallery-card .grid {min-height: 150px;}
+  .plot-card, .history-card {padding: 0.7rem;}
+}
+"""
+
+
+def _prepare_gallery_images(samples: np.ndarray) -> List[np.ndarray]:
+    """Convert normalized grayscale samples to RGB arrays for display."""
+    clipped = np.clip(samples, 0.0, 1.0)
+    uint8_imgs = (clipped * 255).astype(np.uint8)
+    if uint8_imgs.ndim == 3:
+        uint8_imgs = uint8_imgs[..., np.newaxis]
+    return [np.repeat(img, 3, axis=-1) for img in uint8_imgs]
+
+
+def _make_history_plot(history_frames: np.ndarray) -> plt.Figure:
+    """Render up to 10 frames from a sample trajectory in a single row."""
+    if history_frames.ndim == 4 and history_frames.shape[-1] == 1:
+        history_frames = history_frames[..., 0]
+    total_frames = history_frames.shape[0]
+    n_display = min(10, total_frames)
+    if n_display < 1:
+        raise ValueError("History sequence is empty.")
+    indices = np.linspace(0, total_frames - 1, n_display, dtype=int)
+    selected = history_frames[indices]
+
+    fig, axes = plt.subplots(1, n_display, figsize=(2.2 * n_display, 2.2))
+    if n_display == 1:
+        axes = np.array([axes])
+    for idx, ax in enumerate(np.atleast_1d(axes)):
+        ax.axis("off")
+        ax.imshow(selected[idx], cmap="gray")
+        ax.set_title(f"Step {indices[idx] + 1}", fontsize=8, color="#8a450f", pad=6)
+    fig.tight_layout()
+    return fig
+
+
+def _format_summary(
+    *,
+    integrator_label: str,
+    n_steps: int,
+    history_len: int,
+    churn_params: Optional[dict],
+) -> str:
+    sampler_grid = f"""
+      <div class="summary-grid">
+        <span class="summary-pill integrator">{integrator_label}</span>
+        <span class="summary-pill">Steps <strong>{n_steps}</strong></span>
+        <span class="summary-pill">Samples <strong>{N_SAMPLES}</strong></span>
+        <span class="summary-pill">History <strong>{history_len}</strong></span>
+      </div>
+    """.strip()
+
+    churn_block = ""
+    if churn_params:
+        churn_block = f"""
+        <hr class="summary-divider" />
+        <div class="summary-section">
+          <div class="summary-title">Churning</div>
+          <div class="summary-grid">
+            <span class="summary-pill">Rate <strong>{churn_params['stochastic_churn_rate']:.3f}</strong></span>
+            <span class="summary-pill">Min <strong>{churn_params['churn_min']:.3f}</strong></span>
+            <span class="summary-pill">Max <strong>{churn_params['churn_max']:.3f}</strong></span>
+            <span class="summary-pill">Inflation <strong>{churn_params['noise_inflation_factor']:.4f}</strong></span>
+          </div>
+        </div>
+        """.strip()
+
+    return f"""
+    <div class="summary-card">
+      <div class="summary-section">
+        <div class="summary-title">Sampler</div>
+        {sampler_grid}
+      </div>
+      {churn_block}
+    </div>
+    """.strip()
+
+
+def show_history(evt: gr.SelectData, histories: Optional[List[np.ndarray]]):
+    """Render the trajectory plot for the selected sample."""
+    if histories is None or len(histories) == 0:
+        return gr.update(value=None, visible=False), gr.update(
+            value="Click a digit above to explore its diffusion trajectory.",
+            visible=True,
+        )
+
+    index = 0
+    if evt is not None and evt.index is not None:
+        index = evt.index
+        if isinstance(index, (list, tuple)):
+            index = index[-1]
+
+    if not isinstance(index, (int, np.integer)) or index < 0 or index >= len(histories):
+        return gr.update(value=None, visible=False), gr.update(
+            value="Click a digit above to explore its diffusion trajectory.",
+            visible=True,
+        )
+    if histories[index] is None:
+        return gr.update(value=None, visible=False), gr.update(
+            value="Click a digit above to explore its diffusion trajectory.",
+            visible=True,
+        )
+    figure = _make_history_plot(histories[index])
+    return gr.update(value=figure, visible=True), gr.update(visible=False)
+
+
+def generate(
+    integrator_label: str,
+    n_steps: int,
+    seed: int,
+    enable_churn: bool,
+    churn_rate: float,
+    churn_min_value: float,
+    churn_max_value: float,
+    noise_inflation_value: float,
+):
+    """Run sampling with the requested configuration and return UI artifacts."""
+    _, integrator_cfg = resolve_integrator(integrator_label)
+
+    n_steps = int(n_steps)
+    seed = int(seed)
+
+    if not (1 <= n_steps <= MAX_STEPS):
+        raise gr.Error(f"Number of steps must be between 1 and {MAX_STEPS}.")
+
+    supports_churn = integrator_cfg.get("supports_churn", False)
+    churn_params = None
+
+    if enable_churn:
+        if not supports_churn:
+            raise gr.Error("Stochastic churning is only available for deterministic integrators.")
+
+        churn_rate = float(churn_rate)
+        churn_min_value = float(churn_min_value)
+        churn_max_value = float(churn_max_value)
+        noise_inflation_value = float(noise_inflation_value)
+
+        if churn_rate < 0 or churn_rate > 1:
+            raise gr.Error("Churn rate must be within [0, 1].")
+        if churn_min_value < 0 or churn_max_value < 0:
+            raise gr.Error("Churn thresholds must be non-negative.")
+        if churn_max_value < churn_min_value:
+            raise gr.Error("Churn max threshold must be greater than or equal to churn min threshold.")
+        if noise_inflation_value < 1.0 or noise_inflation_value > MAX_NOISE_INFLATION:
+            raise gr.Error(f"Noise inflation factor must be within [1.0, {MAX_NOISE_INFLATION}].")
+
+        churn_params = {
+            "stochastic_churn_rate": churn_rate,
+            "churn_min": churn_min_value,
+            "churn_max": churn_max_value,
+            "noise_inflation_factor": noise_inflation_value,
+        }
+
+    denoiser_state, history = sample_batch(
+        integrator_label,
+        n_steps=n_steps,
+        n_samples=N_SAMPLES,
+        seed=seed,
+        keep_history=True,
+        churn_params=churn_params,
+    )
+
+    integrator_state = denoiser_state.integrator_state
+    samples = jax.device_get(integrator_state.position)
+    samples = np.asarray(samples)
+
+    if samples.ndim == 4 and samples.shape[-1] == 1:
+        samples = samples[..., 0]
+
+    # Diffusion models typically output data in [-1, 1]. Rescale to [0, 1].
+    samples = 0.5 * (samples + 1.0)
+    samples = np.clip(samples, 0.0, 1.0)
+
+    gallery_images = _prepare_gallery_images(samples)
+
+    sample_histories: Optional[List[np.ndarray]] = None
+    if history is not None:
+        history_np = jax.device_get(history)
+        history_np = np.asarray(history_np)
+        history_np = 0.5 * (history_np + 1.0)
+        history_np = np.clip(history_np, 0.0, 1.0)
+        sample_histories = [
+            history_np[:, sample_idx]
+            for sample_idx in range(history_np.shape[1])
+        ]
+    if sample_histories is None:
+        sample_histories = []
+
+    history_len = int(history.shape[0]) if history is not None else 0
+
+    summary_html = _format_summary(
+        integrator_label=integrator_cfg["label"],
+        n_steps=n_steps,
+        history_len=history_len,
+        churn_params=churn_params,
+    )
+
+    gallery_update = gr.update(
+        value=gallery_images,
+        visible=True,
+        interactive=True,
+        height=220,
+    )
+    summary_update = gr.update(value=summary_html)
+    history_reset = gr.update(value=None, visible=False)
+    placeholder_update = gr.update(
+        value="Click a digit above to explore its diffusion trajectory.",
+        visible=True,
+    )
+
+    gr.Info(
+        f"Generated {N_SAMPLES} samples with {integrator_cfg['label']} ({n_steps} steps).",
+        duration=3,
+    )
+
+    return gallery_update, summary_update, history_reset, placeholder_update, sample_histories
+
+
+def _handle_churn_toggle(integrator_label: str, enable_churn: bool):
+    """Toggle churn controls visibility/open state based on integrator support."""
+    _, integrator_cfg = resolve_integrator(integrator_label)
+    supports = integrator_cfg.get("supports_churn", False)
+    enable_effective = supports and enable_churn
+    column_update = gr.update(visible=enable_effective)
+    accordion_update = gr.update(open=enable_effective)
+    return column_update, accordion_update
+
+
+def _handle_integrator_change(integrator_label: str, enable_churn: bool):
+    """Adjust checkbox interactivity and churn panel visibility when integrator changes."""
+    _, integrator_cfg = resolve_integrator(integrator_label)
+    supports = integrator_cfg.get("supports_churn", False)
+    effective_enable = enable_churn if supports else False
+    checkbox_update = gr.update(
+        interactive=supports,
+        value=effective_enable,
+    )
+    column_update, accordion_update = _handle_churn_toggle(integrator_label, effective_enable)
+    return checkbox_update, column_update, accordion_update
+
+
+def _sync_churn_max(churn_min_value: float, current_max_value: float):
+    """Ensure churn_max stays >= churn_min when churn_min changes."""
+    churn_min_value = float(churn_min_value)
+    current_max_value = float(current_max_value)
+    adjusted_max = current_max_value if current_max_value >= churn_min_value else churn_min_value
+    return gr.update(value=adjusted_max)
+
+
+def _sync_churn_min(churn_max_value: float, current_min_value: float):
+    """Ensure churn_min stays <= churn_max when churn_max changes."""
+    churn_max_value = float(churn_max_value)
+    current_min_value = float(current_min_value)
+    adjusted_min = current_min_value if current_min_value <= churn_max_value else churn_max_value
+    return gr.update(value=adjusted_min)
+
+
+def build_ui() -> gr.Blocks:
+    """Create the Gradio Blocks interface."""
+    available_labels = [spec["label"] for spec in INTEGRATORS.values()]
+    default_label = INTEGRATORS["ddim"]["label"]
+
+    with gr.Blocks(
+        title="Diffuse Integrator Explorer",
+        css=CUSTOM_CSS,
+        theme=gr.themes.Soft(primary_hue="orange", secondary_hue="orange"),
+    ) as demo:
+        with gr.Row(elem_id="hero"):
+            gr.Image(
+                value=LOGO_VALUE,
+                show_label=False,
+                interactive=False,
+                elem_classes="hero-logo",
+            )
+            gr.Markdown(
+                """
+                ### Diffuse Integrator Explorer
+                Experiment with deterministic or stochastic samplers from the
+                `diffuse-jax` library. Adjust the number of diffusion steps,
+                hit **Generate Samples**, and compare the five digits rendered
+                in the panel on the right.
+                """.strip(),
+                elem_classes="hero-copy",
+            )
+
+        with gr.Row():
+            with gr.Column(elem_classes="control-card"):
+                gr.Markdown("#### Sampling Controls", elem_classes="control-heading")
+                integrator_input = gr.Dropdown(
+                    choices=available_labels,
+                    value=default_label,
+                    label="Integrator",
+                )
+                steps_input = gr.Slider(
+                    minimum=1,
+                    maximum=MAX_STEPS,
+                    value=DEFAULT_STEPS,
+                    step=1,
+                    label="Number of steps",
+                )
+                seed_input = gr.Number(
+                    value=0,
+                    precision=0,
+                    label="Random seed",
+                    info="Use a different seed to explore new digits.",
+                )
+                with gr.Accordion("Churning controls", open=False, elem_classes="accordion-card") as churn_accordion:
+                    churn_checkbox = gr.Checkbox(
+                        value=False,
+                        label="Enable stochastic churning",
+                        info="Add controlled noise for deterministic integrators.",
+                    )
+                    with gr.Column(visible=False, elem_classes="churn-card") as churn_column:
+                        gr.Markdown(
+                            "**Churning parameters** · tweak how strongly noise is injected during sampling.",
+                            elem_classes="churn-title",
+                        )
+                        churn_rate_input = gr.Slider(
+                            minimum=0.0,
+                            maximum=1.0,
+                            value=DEFAULT_CHURN_RATE,
+                            step=0.01,
+                            label="Churn rate",
+                        )
+                        churn_min_input = gr.Slider(
+                            minimum=0.0,
+                            maximum=1.0,
+                            value=DEFAULT_CHURN_MIN,
+                            step=0.01,
+                            label="Churn min threshold",
+                        )
+                        churn_max_input = gr.Slider(
+                            minimum=0.0,
+                            maximum=1.0,
+                            value=DEFAULT_CHURN_MAX,
+                            step=0.01,
+                            label="Churn max threshold",
+                        )
+                        noise_inflation_input = gr.Slider(
+                            minimum=1.0,
+                            maximum=MAX_NOISE_INFLATION,
+                            value=DEFAULT_NOISE_INFLATION,
+                            step=0.001,
+                            label="Noise inflation factor",
+                        )
+                generate_button = gr.Button(
+                    "Generate Samples",
+                    variant="primary",
+                    elem_classes="generate-button",
+                )
+
+            with gr.Column():
+                details = gr.HTML(
+                    SUMMARY_PLACEHOLDER_HTML,
+                    elem_classes="details-card",
+                    container=False,
+                )
+                gr.Markdown("#### Generated Digit Strip", elem_classes="plot-title")
+                digit_strip = gr.Gallery(
+                    columns=5,
+                    allow_preview=False,
+                    show_fullscreen_button=False,
+                    object_fit="contain",
+                    rows=1,
+                    height=220,
+                    show_label=False,
+                    interactive=True,
+                    elem_classes="gallery-card",
+                    value=[],
+                    container=False,
+                    visible=False,
+                )
+                gr.Markdown("#### Sample Trajectory", elem_classes="plot-title")
+                history_plot = gr.Plot(elem_classes="history-card", show_label=False, visible=False)
+                history_placeholder = gr.Markdown(
+                    "Generate samples, then click a digit above to explore its diffusion trajectory.",
+                    elem_classes="history-placeholder",
+                    visible=True,
+                    container=False,
+                )
+
+        histories_state = gr.State([])
+
+        integrator_input.change(
+            fn=_handle_integrator_change,
+            inputs=[integrator_input, churn_checkbox],
+            outputs=[churn_checkbox, churn_column, churn_accordion],
+        )
+        churn_checkbox.change(
+            fn=_handle_churn_toggle,
+            inputs=[integrator_input, churn_checkbox],
+            outputs=[churn_column, churn_accordion],
+        )
+        churn_min_input.change(
+            fn=_sync_churn_max,
+            inputs=[churn_min_input, churn_max_input],
+            outputs=churn_max_input,
+        )
+        churn_max_input.change(
+            fn=_sync_churn_min,
+            inputs=[churn_max_input, churn_min_input],
+            outputs=churn_min_input,
+        )
+        generate_button.click(
+            fn=generate,
+            inputs=[
+                integrator_input,
+                steps_input,
+                seed_input,
+                churn_checkbox,
+                churn_rate_input,
+                churn_min_input,
+                churn_max_input,
+                noise_inflation_input,
+            ],
+            outputs=[digit_strip, details, history_plot, history_placeholder, histories_state],
+        )
+
+        digit_strip.select(
+            fn=show_history,
+            inputs=[histories_state],
+            outputs=[history_plot, history_placeholder],
+        )
+
+    return demo
+
+
+load_pipeline_assets()
+
+
+if __name__ == "__main__":
+    demo = build_ui()
+    demo.queue().launch()

pipeline.py

@@ -0,0 +1,158 @@
+import functools
+import importlib.util
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+import jax
+from flax import nnx, serialization
+from huggingface_hub import hf_hub_download
+
+from diffuse.diffusion.sde import Flow
+from diffuse.integrator.deterministic import DDIMIntegrator, DPMpp2sIntegrator, EulerIntegrator, HeunIntegrator
+from diffuse.integrator.stochastic import EulerMaruyamaIntegrator
+from diffuse.predictor import Predictor
+from diffuse.timer.base import VpTimer
+from diffuse.denoisers.denoiser import Denoiser
+
+
+@dataclass(frozen=True)
+class PipelineAssets:
+    """Container holding the preloaded model artifacts."""
+
+    model: Any
+    flow: Flow
+    predictor: Predictor
+    x0_shape: Tuple[int, int, int]
+
+
+@functools.lru_cache(maxsize=1)
+def load_pipeline_assets() -> PipelineAssets:
+    """Download the HF model and build the predictor stack once."""
+    model_path = hf_hub_download(repo_id="jcopo/mnist", filename="model.msgpack")
+    config_path = hf_hub_download(repo_id="jcopo/mnist", filename="config.py")
+
+    spec = importlib.util.spec_from_file_location("model_config", config_path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError("Unable to load model config from Hugging Face hub.")
+    config_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(config_module)
+
+    model = config_module.model
+
+    with open(model_path, "rb") as f:
+        state_dict = serialization.from_bytes(None, f.read())
+
+    graphdef, state = nnx.split(model)
+    state.replace_by_pure_dict(state_dict)
+    model = nnx.merge(graphdef, state)
+    model.eval()
+
+    flow = Flow(tf=1.0)
+    predictor = Predictor(
+        model=flow,
+        network=lambda x, t: model(x, t).output,
+        prediction_type="velocity",
+    )
+
+    return PipelineAssets(
+        model=model,
+        flow=flow,
+        predictor=predictor,
+        x0_shape=(28, 28, 1),
+    )
+
+
+INTEGRATORS: Dict[str, Dict[str, Any]] = {
+    "ddim": {
+        "label": "DDIM (Deterministic)",
+        "cls": DDIMIntegrator,
+        "description": "Deterministic DDIM sampler.",
+        "supports_churn": True,
+    },
+    "heun": {
+        "label": "Heun (Deterministic 2nd order)",
+        "cls": HeunIntegrator,
+        "description": "Second-order deterministic integrator.",
+        "supports_churn": True,
+    },
+    "euler": {
+        "label": "Euler (Deterministic)",
+        "cls": EulerIntegrator,
+        "description": "Forward Euler integrator.",
+        "supports_churn": True,
+    },
+    "dpmpp2s": {
+        "label": "DPM++ 2S (Deterministic multi-step)",
+        "cls": DPMpp2sIntegrator,
+        "description": "Deterministic multi-step sampler with second-order accuracy.",
+        "supports_churn": True,
+    },
+    "euler_maruyama": {
+        "label": "Euler-Maruyama (Stochastic)",
+        "cls": EulerMaruyamaIntegrator,
+        "description": "Stochastic sampler with noise at each diffusion step.",
+        "supports_churn": False,
+    },
+}
+
+LABEL_TO_KEY = {spec["label"]: key for key, spec in INTEGRATORS.items()}
+
+
+def resolve_integrator(identifier: str) -> Tuple[str, Dict[str, Any]]:
+    """Resolve either an integrator key or display label to the configuration dict."""
+    if identifier in INTEGRATORS:
+        return identifier, INTEGRATORS[identifier]
+    if identifier in LABEL_TO_KEY:
+        key = LABEL_TO_KEY[identifier]
+        return key, INTEGRATORS[key]
+    raise KeyError(f"Unknown integrator identifier: {identifier}")
+
+
+def build_denoiser(
+    integrator_key: str,
+    n_steps: int,
+    *,
+    churn_params: Optional[Dict[str, float]] = None,
+) -> Denoiser:
+    """Instantiate a denoiser wired with the requested integrator and timer."""
+    if n_steps < 1:
+        raise ValueError("n_steps must be >= 1")
+
+    assets = load_pipeline_assets()
+    _, integrator_cfg = resolve_integrator(integrator_key)
+
+    timer = VpTimer(n_steps=n_steps, eps=0.001, tf=1.0)
+    integrator_kwargs: Dict[str, float] = {}
+    if churn_params:
+        if not integrator_cfg.get("supports_churn", False):
+            raise ValueError(f"Integrator '{integrator_cfg['label']}' does not support stochastic churning.")
+        integrator_kwargs = churn_params
+
+    integrator = integrator_cfg["cls"](model=assets.flow, timer=timer, **integrator_kwargs)
+    return Denoiser(
+        integrator=integrator,
+        model=assets.flow,
+        predictor=assets.predictor,
+        x0_shape=assets.x0_shape,
+    )
+
+
+def sample_batch(
+    integrator_identifier: str,
+    *,
+    n_steps: int,
+    n_samples: int,
+    seed: int,
+    keep_history: bool = False,
+    churn_params: Optional[Dict[str, float]] = None,
+):
+    """Generate a batch of samples for the requested integrator."""
+    if n_samples < 1:
+        raise ValueError("n_samples must be >= 1")
+
+    denoiser = build_denoiser(integrator_identifier, n_steps, churn_params=churn_params)
+    key = jax.random.PRNGKey(seed)
+
+    # The denoiser expects the number of steps to match the timer configuration.
+    state, history = denoiser.generate(key, n_steps, n_samples, keep_history=keep_history)
+    return state, history

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+diffuse-jax
+jax[cpu]
+flax<0.12
+huggingface_hub
+matplotlib
+numpy

script.py

@@ -0,0 +1,22 @@
+from pipeline import load_pipeline_assets, sample_batch
+
+
+def main() -> None:
+    """Example script showing how to invoke the diffusion pipeline."""
+    load_pipeline_assets()
+    print("✅ Model loaded successfully!")
+
+    denoiser_state, history = sample_batch(
+        "ddim",
+        n_steps=10,
+        n_samples=5,
+        seed=456,
+        keep_history=True,
+    )
+
+    print(f"state.position shape: {denoiser_state.integrator_state.position.shape}")
+    print(f"History length: {len(history)}")
+
+
+if __name__ == "__main__":
+    main()