Update app.py

app.py

@@ -8,10 +8,6 @@ from pipeline import StableDiffusionControlLoraV3Pipeline
 from PIL import Image
 import os
 from huggingface_hub import login
-import huggingface_hub
-
-# Initialize ZeroGPU environment
-huggingface_hub.initialize_zerogpu()
 
 # Login using the token
 login(token=os.environ.get("HF_TOKEN"))
@@ -20,11 +16,18 @@ 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
+    torch_dtype=dtype,
+    device_map="auto"  # Let the model handle device placement
 )
 
 # Add conditioning with ow-gbi-control-lora
@@ -34,7 +37,8 @@ unet = unet.add_extra_conditions("ow-gbi-control-lora")
 pipe = StableDiffusionControlLoraV3Pipeline.from_pretrained(
     base_model, 
     unet=unet,
-    torch_dtype=dtype
+    torch_dtype=dtype,
+    device_map="auto"  # Let the model handle device placement
 )
 
 # Use a faster scheduler