Update src/flux/util.py

src/flux/util.py

@@ -306,7 +306,12 @@ def load_flow_model(name: str, device: str | torch.device = "cuda", hf_download:
         and configs[name].repo_flow is not None
         and hf_download
     ):
-        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
+        # ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
+        ckpt_path = hf_hub_download(
+        repo_id="Boese0601/ByteMorpher",
+        filename="dit.safetensors",
+        use_auth_token=os.getenv("HF_TOKEN")
+    )
 
     with torch.device("meta" if ckpt_path is not None else device):
         model = Flux(configs[name].params).to(torch.bfloat16)