Allow setting weight saving flag on hidream_o1

2026-07-16 07:40:25 -06:00
parent 741aeb9ce0
commit 5cb54ba9cc
1 changed files with 1 additions and 1 deletions
--- a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py
+++ b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py
@@ -140,7 +140,7 @@ class HidreamO1Model(BaseModel):
        print(f"Using noise scale: {self.noise_scale}")
        global _GLOBAL_NOISE_SCALE
        _GLOBAL_NOISE_SCALE = self.noise_scale
-        self.is_comfy_weight = False  # save as single file if true
+        self.is_comfy_weight = self.model_config.model_kwargs.get("is_comfy_weight", False)

    # static method to get the noise scheduler
    @staticmethod