Update model_loading.py

model_loading.py

@@ -46,8 +46,10 @@ def import_model(modelID, modelType):
         elif modelType == 'FluxPipeline':
             from diffusers import FluxPipeline
             T2IModel = FluxPipeline.from_pretrained(modelID, torch_dtype=torch.bfloat16)
+            T2IModel.enable_model_cpu_offload()
         else:
             from diffusers import AutoPipelineForText2Image
             T2IModel = AutoPipelineForText2Image.from_pretrained(modelID, torch_dtype=torch.float16)
+            
         T2IModel.to("cuda")
     return T2IModel