Update main.py

main.py

@@ -1,7 +1,8 @@
 run_api = False
-SSD_1B = False
+is_ssd = False
+is_sdxl = True
+is_sdxl_turbo=False
 import os
-
 # Use GPU
 gpu_info = os.popen("nvidia-smi").read()
 if "failed" in gpu_info:
@@ -47,8 +48,6 @@ import PIL
 import base64
 import io
 import torch
-from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
-
 # SDXL
 from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
 
@@ -67,51 +66,52 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "default_secret")
 if is_gpu:
     # Uncomment the following line if you want to enable CUDA launch blocking
     os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
-    torch_dtype=torch.float16
-    variant="fp16"
 else:
     # Uncomment the following line if you want to use CPU instead of GPU
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    torch_dtype=torch.float32
-    variant="fp32"
-
-
 # Get the current directory
 current_dir = os.getcwd()
 model_path = os.path.join(current_dir)
-
 # Set the cache path
 cache_path = os.path.join(current_dir, "cache")
 
-if not SSD_1B:
-
+def load_pipeline(use_cuda):
+    device = "cuda" if use_cuda and torch.cuda.is_available() else "cpu"
+    if device == "cuda":
+        torch.cuda.max_memory_allocated(device=device)
+        torch.cuda.empty_cache()
+        pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+        pipe.enable_xformers_memory_efficient_attention()
+        pipe = pipe.to(device)
+        torch.cuda.empty_cache()
+    else:
+        pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
+        pipe = pipe.to(device)
+    return pipe
+
+if is_sdxl:
+    torch_dtype=torch.float16
+    variant="fp16"
     unet = UNet2DConditionModel.from_pretrained(
         "latent-consistency/lcm-sdxl",        
         torch_dtype=torch_dtype,
         variant=variant,
         cache_dir=cache_path,
     )
-
-    #  model_id="stabilityai/stable-diffusion-xl-base-1.0"
-    model_id="stabilityai/sdxl-turbo"
-    #pipe = DiffusionPipeline.from_pretrained(
-    #    model_id=model_id,
-    #    unet=unet,
-    #    torch_dtype=torch_dtype,
-    #    variant=variant,
-    #    cache_dir=cache_path,
-   # )
-    from diffusers import StableDiffusionPipeline
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
-
-
+    model_id="stabilityai/stable-diffusion-xl-base-1.0"
+    pipe = DiffusionPipeline.from_pretrained(
+        model_id=model_id,
+        unet=unet,
+        torch_dtype=torch_dtype,
+        variant=variant,
+        cache_dir=cache_path,
+    )
     pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
     if torch.cuda.is_available():
         pipe.to("cuda")
-else:
+if is_ssd:
     # SSD-1B
     from diffusers import LCMScheduler, AutoPipelineForText2Image
-
     pipe = AutoPipelineForText2Image.from_pretrained(
         "segmind/SSD-1B",
         torch_dtype=torch.float16,
@@ -121,11 +121,13 @@ else:
     pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
     if torch.cuda.is_available():
         pipe.to("cuda")
-
     # load and fuse
     pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
     pipe.fuse_lora()
 
+if is_sdxl_turbo:
+    use_cuda=is_gpu
+    pipe = load_pipeline(use_cuda)
 
 def generate(
     prompt: str,