Update src/flux/pipeline_tools.py

src/flux/pipeline_tools.py

@@ -515,13 +515,13 @@ class CustomFluxPipeline:
         if config["model"].get("dit_quant", "None")=="int8-quanto":
             self.pipe = FluxPipeline.from_pretrained("diffusers/FLUX.1-dev-torchao-int8", 
                                                     torch_dtype=torch_dtype,
-                                                    use_safetensors=False).to(device)
+                                                    use_safetensors=False, device_map=device)
 
         else:
             self.pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", 
                                                     torch_dtype=torch_dtype).to(device)
         
-        self.pipe.enable_sequential_cpu_offload()
+        # self.pipe.enable_sequential_cpu_offload()
 
         self.config = config
         self.device = device