Add files using upload-large-folder tool

README.md

@@ -49,6 +49,7 @@ import torch
 
 pipe = DiffusionPipeline.from_pretrained(
     "BiliSakura/AeroGen",
+    custom_pipeline="pipeline.py",
     trust_remote_code=True,
 )
 pipe = pipe.to("cuda")

model_index.json

@@ -1,5 +1,5 @@
 {
-  "_class_name": ["pipeline", "AeroGenPipeline"],
+  "_class_name": "AeroGenPipeline",
   "_diffusers_version": "0.25.0",
   "condition_encoder": [
     "pipeline",

modular_pipeline.py

@@ -106,12 +106,33 @@ def load_component(model_path: Path, name: str):
     with open(comp_path / "config.json") as f:
         cfg = json.load(f)
 
-    # Diffusers native format (e.g. AutoencoderKL.save_pretrained): no "target" key
-    if "target" not in cfg and name == "vae":
+    # VAE loading: support both native diffusers format and legacy LDM config.
+    if name == "vae" and "target" not in cfg:
         from diffusers import AutoencoderKL
         return AutoencoderKL.from_pretrained(comp_path)
 
-    component = _instantiate_from_config(cfg)
+    if name == "vae" and cfg.get("target") == "ldm.models.autoencoder.AutoencoderKL":
+        from diffusers import AutoencoderKL
+
+        ddconfig = (cfg.get("params") or {}).get("ddconfig") or {}
+        ch = int(ddconfig.get("ch", 128))
+        ch_mult = ddconfig.get("ch_mult") or [1, 2, 4, 4]
+        block_out_channels = [ch * int(m) for m in ch_mult]
+
+        component = AutoencoderKL(
+            in_channels=int(ddconfig.get("in_channels", 3)),
+            out_channels=int(ddconfig.get("out_ch", 3)),
+            down_block_types=["DownEncoderBlock2D"] * len(block_out_channels),
+            up_block_types=["UpDecoderBlock2D"] * len(block_out_channels),
+            block_out_channels=block_out_channels,
+            layers_per_block=int(ddconfig.get("num_res_blocks", 2)),
+            latent_channels=int(ddconfig.get("z_channels", 4)),
+            sample_size=int(ddconfig.get("resolution", 256)),
+            act_fn="silu",
+            norm_num_groups=32,
+        )
+    else:
+        component = _instantiate_from_config(cfg)
     safetensors_path = comp_path / "diffusion_pytorch_model.safetensors"
     bin_path = comp_path / "diffusion_pytorch_model.bin"
     if safetensors_path.exists():
@@ -131,7 +152,13 @@ def load_component(model_path: Path, name: str):
     # Older checkpoints may have been saved without the "model." prefix.
     if name == "unet" and state and not any(k.startswith("model.") for k in state.keys()):
         state = {"model." + k: v for k, v in state.items()}
-    component.load_state_dict(state, strict=True)
+    if name == "vae" and cfg.get("target") == "ldm.models.autoencoder.AutoencoderKL":
+        from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint
+
+        state = convert_ldm_vae_checkpoint(state, dict(component.config))
+        component.load_state_dict(state, strict=False)
+    else:
+        component.load_state_dict(state, strict=True)
     component.eval()
     return component
 

pipeline.py

@@ -36,14 +36,27 @@ from diffusers import DDIMScheduler, DiffusionPipeline
 from diffusers.utils import BaseOutput
 from PIL import Image
 
-from modular_pipeline import (
-    ensure_ldm_path,
-    ensure_ldm_path_from_config,
-    load_component,
-    load_components,
-    create_scheduler,
-    _instantiate_from_config,
-)
+try:
+    # Dynamic modules loaded by diffusers are executed as package modules.
+    from .modular_pipeline import (
+        ensure_ldm_path,
+        ensure_ldm_path_from_config,
+        load_component,
+        load_components,
+        create_scheduler,
+        _instantiate_from_config,
+    )
+except ImportError:
+    # Fallback for direct local execution (e.g. `python pipeline.py`).
+    import importlib
+
+    _mp = importlib.import_module("modular_pipeline")
+    ensure_ldm_path = _mp.ensure_ldm_path
+    ensure_ldm_path_from_config = _mp.ensure_ldm_path_from_config
+    load_component = _mp.load_component
+    load_components = _mp.load_components
+    create_scheduler = _mp.create_scheduler
+    _instantiate_from_config = _mp._instantiate_from_config
 
 
 @dataclass