Update ssh

model/models/huggingface_models.py

@@ -6,20 +6,20 @@ import torch
 
 
 def load_huggingface_model(model_name, model_type):
-    # if model_name == "SD-turbo":
-    #     pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo", torch_dtype=torch.float16, variant="fp16")
-    # elif model_name == "SDXL-turbo":
-    #     pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
-    # else:
-    #     raise NotImplementedError
-    # pipe = pipe.to("cuda")
     if model_name == "SD-turbo":
-        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo")
+        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo", torch_dtype=torch.float16, variant="fp16")
     elif model_name == "SDXL-turbo":
-        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
+        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
     else:
         raise NotImplementedError
-    pipe = pipe.to("cpu")
+    pipe = pipe.to("cuda")
+    # if model_name == "SD-turbo":
+    #     pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo")
+    # elif model_name == "SDXL-turbo":
+    #     pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
+    # else:
+    #     raise NotImplementedError
+    # pipe = pipe.to("cpu")
     return pipe