Add files using upload-large-folder tool

README.md

@@ -9,6 +9,11 @@ tags:
   - latent-diffusion
   - custom-pipeline
   - arxiv:2411.16969
+widget:
+  - src: demo_images/input.jpeg
+    prompt: BRCA sample conditioned on demo SSL feature (mag=0)
+    output:
+      url: demo_images/output.jpeg
 ---
 
 # BiliSakura/ZoomLDM-brca
@@ -60,6 +65,20 @@ out = pipe(
 images = out.images
 ```
 
+## Demo Generation (dataset-backed)
+
+This repo includes `run_demo_inference.py`, which uses local repo assets only:
+
+- image: `demo_images/input.jpeg`
+- SSL feature: `demo_data/0_ssl_feat.npy`
+- magnification label: `0`
+
+Run:
+
+```bash
+python run_demo_inference.py
+```
+
 ## Limitations
 
 - Requires correctly precomputed BRCA conditioning features.

demo_data/0_ssl_feat.npy

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

pipeline_zoomldm.py

@@ -44,6 +44,10 @@ def _ensure_local_ldm_on_path():
 
 
 _ensure_local_ldm_on_path()
+# Register module alias so diffusers component loading can resolve
+# model_index entries like "pipeline_zoomldm" even when this file is loaded
+# under a dynamic module name (e.g. diffusers_modules.local.*).
+sys.modules["pipeline_zoomldm"] = sys.modules[__name__]
 
 
 def _get_class(target: str):
@@ -299,6 +303,17 @@ class ZoomLDMPipeline(DiffusionPipeline):
             path = Path(snapshot_download(pretrained_model_name_or_path))
 
         path = path.resolve()
+        component_names = {"unet", "vae", "conditioning_encoder"}
+        # When diffusers loads components, it may call this class with a path like ".../unet".
+        requested_component = None
+        if path.name in component_names and (path / "config.json").exists():
+            requested_component = path.name
+            path = path.parent
+
+        # Also support explicit component requests via subfolder.
+        subfolder = kwargs.pop("subfolder", None)
+        if requested_component is None and subfolder in component_names:
+            requested_component = subfolder
 
         def _is_diffusers_model_dir(candidate: Path) -> bool:
             required = [
@@ -333,8 +348,6 @@ class ZoomLDMPipeline(DiffusionPipeline):
                 )
             model_dir = candidate_dirs[0]
 
-        scheduler = DDIMScheduler.from_pretrained(model_dir / "scheduler")
-
         _TARGETS = {
             "unet": "ldm.modules.diffusionmodules.openaimodel.UNetModel",
             "vae": "ldm.models.autoencoder.VQModelInterface",
@@ -381,6 +394,12 @@ class ZoomLDMPipeline(DiffusionPipeline):
             component.eval()
             return component
 
+        # Diffusers component-loading path: return a single module.
+        if requested_component is not None:
+            return load_custom_component(requested_component)
+
+        scheduler = DDIMScheduler.from_pretrained(model_dir / "scheduler")
+
         unet = load_custom_component("unet")
         vae = load_custom_component("vae")
         conditioning_encoder = load_custom_component("conditioning_encoder")

run_demo_inference.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""Run ZoomLDM-BRCA demo inference using local demo assets."""
+
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers import DiffusionPipeline
+
+
+def preprocess_brca_ssl(npy_path: Path) -> torch.Tensor:
+    # Copied from dataset material:
+    # 1) cast to float32
+    # 2) normalize per-feature
+    # 3) reshape to (1024, h, h)
+    # 4) adaptive avg pool to max 8x8 if needed
+    feat = np.load(npy_path).astype(np.float32)
+    if feat.ndim == 1:
+        feat = feat[:, None]
+    mean = feat.mean(axis=0, keepdims=True)
+    std = feat.std(axis=0, keepdims=True)
+    feat = (feat - mean) / (std + 1e-8)
+    h = int(np.sqrt(feat.shape[1]))
+    feat = torch.tensor(feat.reshape((-1, h, h))).float()  # (1024, h, h)
+    if h > 8:
+        feat = F.adaptive_avg_pool2d(feat, (8, 8))
+    return feat
+
+
+def main() -> None:
+    repo = Path(__file__).resolve().parent
+    demo_dir = repo / "demo_images"
+    demo_data = repo / "demo_data"
+    demo_dir.mkdir(exist_ok=True)
+
+    # Use repo-local demo assets only.
+    src_img = demo_dir / "input.jpeg"
+    src_feat = demo_data / "0_ssl_feat.npy"
+    if not src_img.exists():
+        raise FileNotFoundError(f"Missing demo input image: {src_img}")
+    if not src_feat.exists():
+        raise FileNotFoundError(f"Missing demo SSL feature: {src_feat}")
+
+    ssl_feat = preprocess_brca_ssl(src_feat).unsqueeze(0).to("cuda")  # (1, 1024, h, h)
+    magnification = torch.tensor([0], device="cuda", dtype=torch.long)
+
+    pipe = DiffusionPipeline.from_pretrained(
+        str(repo),
+        custom_pipeline=str(repo / "pipeline_zoomldm.py"),
+        trust_remote_code=True,
+        local_files_only=True,
+    ).to("cuda")
+
+    out = pipe(
+        ssl_features=ssl_feat,
+        magnification=magnification,
+        num_inference_steps=50,
+        guidance_scale=2.0,
+        generator=torch.Generator(device="cuda").manual_seed(42),
+    )
+    out.images[0].save(demo_dir / "output.jpeg")
+    print(f"Saved {demo_dir / 'input.jpeg'}")
+    print(f"Saved {demo_dir / 'output.jpeg'}")
+
+
+if __name__ == "__main__":
+    main()