Adjust buckets for flex2

2025-04-02 06:47:41 -06:00
parent 3d131fb27a
commit a42c5a1de5
2 changed files with 4 additions and 1 deletions
--- a/extensions_built_in/flex2/flex2.py
+++ b/extensions_built_in/flex2/flex2.py
@@ -72,6 +72,9 @@ class Flex2(BaseModel):
    def get_train_scheduler():
        return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)

+    def get_bucket_divisibility(self):
+        return 16
+
    def load_model(self):
        dtype = self.torch_dtype
        self.print_and_status_update("Loading Flux2 model")