fix: add @spaces.GPU decorators, dynamic device/dtype fallback, and Gradio 6.0 launch params

README.md

@@ -5,6 +5,7 @@ colorFrom: gray
 colorTo: indigo
 sdk: gradio
 sdk_version: 6.20.0
+python_version: "3.10"
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -6,8 +6,10 @@ import numpy as np
 import requests
 import torch
 import gc
+import spaces
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 # Download and Create SAM Model
 
@@ -34,11 +36,11 @@ gc.collect()
 print("Creating ControlNet Pipeline")
 
 controlnet = ControlNetModel.from_pretrained(
-    "mfidabel/controlnet-segment-anything", torch_dtype=torch.float16
+    "mfidabel/controlnet-segment-anything", torch_dtype=dtype
 ).to(device)
 
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, safety_check=None
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=dtype
 ).to(device)
 
 
@@ -109,6 +111,7 @@ def show_anns(anns):
     
     return final_img
 
+@spaces.GPU(duration=30)
 def segment_image(image, seed = 0):
     # Generate Masks
     np.random.seed(int(seed))
@@ -121,6 +124,7 @@ def segment_image(image, seed = 0):
     torch.cuda.empty_cache()
     return map
 
+@spaces.GPU(duration=60)
 def infer(prompts, negative_prompts, image, num_inference_steps = 50, seed = 4, num_samples = 4):
     try:
         # Segment Image
@@ -161,7 +165,7 @@ prompt = gr.Textbox(lines=1, label="Prompt", value=default_example[0])
 negative_prompt = gr.Textbox(lines=1, label="Negative Prompt", value=default_example[1])
 
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             # Title
@@ -221,4 +225,4 @@ with gr.Blocks(css=css) as demo:
                      outputs=segm_img)
     
 demo.queue()
-demo.launch()
+demo.launch(css=css)