Update app.py

app.py

@@ -38,6 +38,11 @@ ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
 GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
 
 LORA_REPO = os.environ.get("H3_LORA_REPO", "dagloop5/LoRA")
+# 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.
+CUSTOM_TRANSFORMER_REPO = os.environ.get("H3_CUSTOM_TRANSFORMER_REPO", "TenStrip/10Eros-Max")
+CUSTOM_TRANSFORMER_FILE = os.environ.get("H3_CUSTOM_TRANSFORMER_FILE", "10Eros_Max_h3_TURBO-hybrid_beta3.safetensors")
 # 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.
 LORA_FILES = {
@@ -388,6 +393,34 @@ 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 yet, then `load_state_dict(assign=True)`
+            # materializes real tensors directly from the finetune's own state dict — the only real allocation
+            # in this path, and it happens exactly once, for the weights actually being kept.
+            #
+            # Setting `pipe.transformer` here, before `load_components()` runs, is what makes `load_components()`
+            # skip fetching it at all: its own `names=None` branch only loads components where
+            # `getattr(self, name, None) is None`.
+            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)
+            # `strict=True`: a finetune with a genuinely different key set is a real architecture mismatch, not
+            # something to load partially and hope for the best on.
+            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)