Update model_loading.py

model_loading.py

@@ -37,12 +37,15 @@ def import_model(modelID, modelType):
         if modelType == 'StableDiffusionXLPipeline':
             from diffusers import StableDiffusionXLPipeline
             T2IModel = StableDiffusionXLPipeline.from_pretrained(modelID, torch_dtype=torch.float16)
+            T2IModel.to("cuda")
         elif modelType == 'LatentConsistencyModelPipeline':
             from diffusers import DiffusionPipeline
             T2IModel = DiffusionPipeline.from_pretrained(modelID, torch_dtype=torch.float16)
+            T2IModel.to("cuda")
         elif modelType == 'StableDiffusion3Pipeline':
             from diffusers import StableDiffusion3Pipeline
             T2IModel = StableDiffusion3Pipeline.from_pretrained(modelID, torch_dtype=torch.bfloat16)
+            T2IModel.to("cuda")
         elif modelType == 'FluxPipeline':
             from diffusers import FluxPipeline
             T2IModel = FluxPipeline.from_pretrained(modelID, torch_dtype=torch.bfloat16)
@@ -50,6 +53,14 @@ def import_model(modelID, modelType):
         else:
             from diffusers import AutoPipelineForText2Image
             T2IModel = AutoPipelineForText2Image.from_pretrained(modelID, torch_dtype=torch.float16)
+            T2IModel.to("cuda")
             
+    if 'StableDiffusionXLPipeline' in modelType.split(','):
+        from diffusers import StableDiffusionXLPipeline
+        T2IModel = StableDiffusionXLPipeline.from_pretrained(modelID, torch_dtype=torch.float16)
         T2IModel.to("cuda")
+    try:
+        T2IModel.safety_checker = None
+    except:
+        pass    # if the model does not contain a safety checker no need to remove it
     return T2IModel