Update app.py

app.py

@@ -1,3 +1,4 @@
+import torch
 import spaces
 import gradio as gr
 
@@ -5,11 +6,13 @@ from diffusers import StableDiffusionPipeline
 
 
 repo = "Tramac/style-portrait-v1-5"
-pipeline = StableDiffusionPipeline.from_pretrained(
-    repo,
-    torch_dtype=torch.float16,
-    variant="fp16"
-).to("cuda")
+# Ensure model and scheduler are initialized in GPU-enabled function
+if torch.cuda.is_available():
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        repo,
+        torch_dtype=torch.float16,
+        variant="fp16"
+    ).to("cuda")
 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, neg_prompt):