Update app.py

app.py

@@ -40,7 +40,7 @@ masterpiece, newest, absurdres
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -72,6 +72,20 @@ def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
     }
     return scheduler_factory_map.get(name, lambda: None)()
 
+def load_pipeline(model_name):
+    if torch.cuda.is_available():
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            custom_pipeline="lpw_stable_diffusion_xl",
+            safety_checker=None,
+            use_safetensors=True,
+            add_watermarker=False,
+            use_auth_token=HF_TOKEN
+        )
+    pipe.to(device)
+    return pipe
+
 @spaces.GPU(enable_queue=False)
 def generate(
     prompt: str,
@@ -81,9 +95,10 @@ def generate(
     height: int = 1024,
     guidance_scale: float = 5.0,
     num_inference_steps: int = 26,
-    sampler: str = "Euler a",
+    sampler: str = "Eul""er a",
     clip_skip: int = 1,
 ):
+    """
     if torch.cuda.is_available():
         pipe = StableDiffusionXLPipeline.from_pretrained(
             MODEL,
@@ -94,6 +109,7 @@ def generate(
             add_watermarker=False,
             use_auth_token=HF_TOKEN
         )
+    """
     
     generator = seed_everything(seed)
     pipe.scheduler = get_scheduler(pipe.scheduler.config, sampler)