Update app.py

app.py

@@ -61,20 +61,6 @@ 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
 def generate(
     prompt: str,
@@ -88,6 +74,17 @@ def generate(
     clip_skip: int = 1,
     progress=gr.Progress(track_tqdm=True),
 ):
+    if torch.cuda.is_available():
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            MODEL,
+            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)
     
     generator = seed_everything(seed)
     pipe.scheduler = get_scheduler(pipe.scheduler.config, sampler)
@@ -116,13 +113,6 @@ def generate(
     except Exception as e:
         print(f"An error occurred: {e}")
 
-if torch.cuda.is_available():
-    pipe = load_pipeline(MODEL)
-    pipe.to(device)
-    print("Loaded on Device!")
-else:
-    pipe = None
-
 with gr.Blocks(
     theme=gr.themes.Soft()
 ) as demo: