Update app.py

app.py

@@ -8,6 +8,7 @@ from pipeline import StableDiffusionControlLoraV3Pipeline
 from PIL import Image
 import os
 from huggingface_hub import login
+import spaces
 
 # Login using the token
 login(token=os.environ.get("HF_TOKEN"))
@@ -16,18 +17,11 @@ login(token=os.environ.get("HF_TOKEN"))
 base_model = "runwayml/stable-diffusion-v1-5"
 dtype = torch.float16  # A100 works better with float16
 
-try:
-    # Check if CUDA is available
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-except:
-    device = "cpu"
-
 # Load the custom UNet
 unet = UNet2DConditionModelEx.from_pretrained(
     base_model,
     subfolder="unet",
-    torch_dtype=dtype,
-    device_map="auto"  # Let the model handle device placement
+    torch_dtype=dtype
 )
 
 # Add conditioning with ow-gbi-control-lora
@@ -37,8 +31,7 @@ unet = unet.add_extra_conditions("ow-gbi-control-lora")
 pipe = StableDiffusionControlLoraV3Pipeline.from_pretrained(
     base_model, 
     unet=unet,
-    torch_dtype=dtype,
-    device_map="auto"  # Let the model handle device placement
+    torch_dtype=dtype
 )
 
 # Use a faster scheduler
@@ -61,6 +54,7 @@ def get_canny_image(image, low_threshold=100, high_threshold=200):
     canny_image = np.stack([canny_image] * 3, axis=-1)
     return Image.fromarray(canny_image)
 
+@spaces.GPU(duration=120)  # Set GPU allocation duration to 120 seconds
 def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps, low_threshold, high_threshold):
     canny_image = get_canny_image(input_image, low_threshold, high_threshold)