Update src/flux/pipeline_tools.py

src/flux/pipeline_tools.py

@@ -510,7 +510,8 @@ class CustomFluxPipeline:
         ckpt_root_condition=None,
         torch_dtype=torch.bfloat16,
     ):
-        model_path = os.getenv("FLUX_MODEL_PATH", "black-forest-labs/FLUX.1-dev")
+        
+        model_path = os.getenv("FLUX_MODEL_PATH", "diffusers/FLUX.1-dev-torchao-int8" if config["model"].get("dit_quant", "None")=="int8-quanto" else "black-forest-labs/FLUX.1-dev")
         print("[CustomFluxPipeline] Loading FLUX Pipeline")
         self.pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch_dtype).to(device)
         self.pipe.enable_sequential_cpu_offload()
@@ -518,8 +519,8 @@ class CustomFluxPipeline:
         self.config = config
         self.device = device
         self.dtype = torch_dtype
-        if config["model"].get("dit_quant", "None") != "None":
-            quantization(self.pipe, config["model"]["dit_quant"])
+        # if config["model"].get("dit_quant", "None") != "None":
+        #     quantization(self.pipe, config["model"]["dit_quant"])
 
         self.modulation_adapters = []
         self.pipe.modulation_adapters = []