Update app.py

app.py

@@ -41,8 +41,10 @@ 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")
+# A proper diffusers-native, sharded transformer checkpoint (its own `config.json` +
+# `diffusion_pytorch_model-*-of-*.safetensors` + index, at the repo root) — 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", "ibyteohdear/10Eros-Max-Transformer")
 # 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 = {
@@ -395,31 +397,19 @@ def load_models() -> str | None:
         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.
+            # A real sharded diffusers repo, not a single file — `from_pretrained` handles shard resolution and
+            # low-memory (meta-device-backed) loading itself, so nothing manual is needed here.
             #
             # 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
-
+            # skip fetching the official transformer entirely: its own `names=None` branch only loads components
+            # where `getattr(self, name, None) is None`.
             from diffusers.models import MiniMaxH3Transformer3DModel
 
-            config, _ = MiniMaxH3Transformer3DModel.load_config(
-                MODEL_REPO, subfolder="transformer", return_unused_kwargs=True
+            custom_transformer = MiniMaxH3Transformer3DModel.from_pretrained(
+                CUSTOM_TRANSFORMER_REPO, torch_dtype=torch.bfloat16
             )
-            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)
+            print(f"[gen] transformer replaced with {CUSTOM_TRANSFORMER_REPO}", flush=True)
 
         pipe.load_components(dtype=torch.bfloat16)
         pipe.transformer.set_attention_backend(ATTENTION)