Update app.py

app.py

@@ -37,6 +37,16 @@ PLACEMENT = os.environ.get("H3_PLACEMENT", "pack").lower()
 ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
 GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
 
+# A finetuned transformer, as a single monolithic safetensors file rather than MODEL_REPO's own sharded
+# `transformer/` subfolder — everything else (VAE, schedulers, config) still comes from MODEL_REPO. Empty by
+# default, which reproduces the official weights exactly. Confirmed diffusers-native key naming (not ComfyUI,
+# not pruned-AdaLN) via `xal2077/PinkCherry_MiniMax-H3-Demo`'s own working `load_state_dict(strict=True)` call
+# against an earlier release of the same lineage.
+CUSTOM_TRANSFORMER_REPO = os.environ.get("H3_CUSTOM_TRANSFORMER_REPO", "coeboy/PinkCherry_MiniMax-H3")
+CUSTOM_TRANSFORMER_FILE = os.environ.get(
+    "H3_CUSTOM_TRANSFORMER_FILE", "v1-final-fl2va/PinkCherry_v1_bf16_fla2va_H3.safetensors"
+)
+
 LORA_REPO = os.environ.get("H3_LORA_REPO", "dagloop5/LoRA")
 # Each entry is (repo, filename) so a LoRA can come from any repo, not just LORA_REPO — the two Lightx2v files
 # live in lightx2v/Minimax-h3-Turbo, not dagloop5/LoRA.
@@ -423,6 +433,27 @@ def load_models() -> str | None:
         blocks = MiniMaxH3GeneratorBlocks()
         print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
         pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
+
+        if CUSTOM_TRANSFORMER_REPO:
+            # `load_config` fetches only `transformer/config.json` (a few KB) — not the 61.7 GiB of weights
+            # `load_components` would otherwise pull from MODEL_REPO. Constructed on `torch.device("meta")` so
+            # the architecture exists with no real memory behind it, then `load_state_dict(assign=True)`
+            # materializes real tensors straight from the finetune's own file — the only real allocation here.
+            from huggingface_hub import hf_hub_download
+            from safetensors.torch import load_file
+
+            from diffusers.models import MiniMaxH3Transformer3DModel
+
+            config, _ = MiniMaxH3Transformer3DModel.load_config(
+                MODEL_REPO, subfolder="transformer", return_unused_kwargs=True
+            )
+            with torch.device("meta"):
+                custom_transformer = MiniMaxH3Transformer3DModel.from_config(config)
+            custom_path = hf_hub_download(CUSTOM_TRANSFORMER_REPO, CUSTOM_TRANSFORMER_FILE)
+            custom_transformer.load_state_dict(load_file(custom_path), strict=True, assign=True)
+            pipe.update_components(transformer=custom_transformer)
+            print(f"[gen] transformer replaced with {CUSTOM_TRANSFORMER_REPO}/{CUSTOM_TRANSFORMER_FILE}", flush=True)
+
         pipe.load_components(dtype=torch.bfloat16)
         pipe.transformer.set_attention_backend(ATTENTION)