test

app.py

@@ -13,6 +13,13 @@ CONTROLNET_FILENAME = "ControlAny-SDXL/anime_canny/controlnet.safetensors"
 CACHE_DIR = None
 
 
+@spaces.GPU
+def optimize_pipeline(pipeline):
+    pipeline.to(dtype=torch.float16)
+    pipeline.enable_xformers_memory_efficient_attention()
+    return pipeline
+
+
 def ui():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model_file = hf_hub_download(
@@ -41,10 +48,9 @@ def ui():
         device=device,
         hf_cache_dir=CACHE_DIR,
         use_safetensors=True,
-        enable_xformers_memory_efficient_attention=True,
+        enable_xformers_memory_efficient_attention=torch.cuda.is_available(),
     )
-    with spaces.GPU:
-        pipeline = pipeline.to(dtype=torch.float16)
+    pipeline = optimize_pipeline(pipeline)
 
     preprocessors = ['canny']
     schedulers = ['Euler A', 'UniPC', 'Euler', 'DDIM', 'DDPM']

utils/tools.py

@@ -112,11 +112,11 @@ def get_pipeline(
 
     pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
     pipeline.set_progress_bar_config()
-    pipeline = pipeline.to(device, dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+    pipeline = pipeline.to(device, dtype=torch.float16)
 
     if lora_path is not None:
         pipeline.load_lora_weights(lora_path)
-    if enable_xformers_memory_efficient_attention and torch.cuda.is_available():
+    if enable_xformers_memory_efficient_attention:
         pipeline.enable_xformers_memory_efficient_attention()
 
     return pipeline