Update for ZeroGPU

app.py

@@ -1,22 +1,24 @@
-import gradio as gr
 import os
 import json
-from openai import OpenAI
-from geopy.geocoders import Nominatim
-from folium import Map, GeoJson
-from gradio_folium import Folium
 import cv2
 import numpy as np
 import torch
-from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
 from PIL import Image
 import io
+import gradio as gr
+from openai import OpenAI
+from geopy.geocoders import Nominatim
+from folium import Map, GeoJson
+from gradio_folium import Folium
+from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
+import spaces
 
 # Initialize APIs
 openai_client = OpenAI(api_key=os.environ['OPENAI_API_KEY'])
 geolocator = Nominatim(user_agent="geoapi")
 
 # Function to fetch coordinates
+@spaces.GPU
 def get_geo_coordinates(location_name):
     try:
         location = geolocator.geocode(location_name)
@@ -28,6 +30,7 @@ def get_geo_coordinates(location_name):
         return None
 
 # Function to process OpenAI chat response
+@spaces.GPU
 def process_openai_response(query):
     response = openai_client.chat.completions.create(
         model="gpt-4o-mini",
@@ -45,17 +48,18 @@ def process_openai_response(query):
     return json.loads(response.choices[0].message.content)
 
 # Generate GeoJSON from OpenAI response
+@spaces.GPU
 def generate_geojson(response):
     feature_type = response['output']['feature_representation']['type']
     city_names = response['output']['feature_representation']['cities']
     properties = response['output']['feature_representation']['properties']
-    
+
     coordinates = []
     for city in city_names:
         coord = get_geo_coordinates(city)
         if coord:
             coordinates.append(coord)
-    
+
     if feature_type == "Polygon":
         coordinates.append(coordinates[0])  # Close the polygon
 
@@ -71,8 +75,8 @@ def generate_geojson(response):
         }]
     }
 
-
 # Function to compute bounds from GeoJSON
+@spaces.GPU
 def get_bounds(geojson):
     coordinates = []
     for feature in geojson["features"]:
@@ -94,6 +98,7 @@ def get_bounds(geojson):
     return [[min(lats), min(lngs)], [max(lats), max(lngs)]]
 
 # Generate map image
+@spaces.GPU
 def save_map_image(geojson_data):
     m = Map()
     geo_layer = GeoJson(geojson_data, name="Feature map")
@@ -110,10 +115,10 @@ controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpai
 pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     "stable-diffusion-v1-5/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16
 )
-#pipeline.enable_model_cpu_offload()
+# ZeroGPU compatibility
 pipeline.to('cuda')
 
-
+@spaces.GPU
 def make_inpaint_condition(init_image, mask_image):
     init_image = np.array(init_image.convert("RGB")).astype(np.float32) / 255.0
     mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0
@@ -124,6 +129,7 @@ def make_inpaint_condition(init_image, mask_image):
     init_image = torch.from_numpy(init_image)
     return init_image
 
+@spaces.GPU
 def generate_satellite_image(init_image_path, mask_image_path, prompt):
     init_image = Image.open(init_image_path)
     mask_image = Image.open(mask_image_path)
@@ -132,24 +138,25 @@ def generate_satellite_image(init_image_path, mask_image_path, prompt):
     return result.images[0]
 
 # Gradio UI
+@spaces.GPU
 def handle_query(query):
     # Process OpenAI response
     response = process_openai_response(query)
     geojson_data = generate_geojson(response)
-    
+
     # Save map image
     map_image_path = save_map_image(geojson_data)
-    
+
     # Generate mask for ControlNet
     empty_map = cv2.imread("empty_map_image.png")
     map_image = cv2.imread(map_image_path)
     difference = cv2.absdiff(cv2.cvtColor(empty_map, cv2.COLOR_BGR2GRAY), cv2.cvtColor(map_image, cv2.COLOR_BGR2GRAY))
     _, mask = cv2.threshold(difference, 15, 255, cv2.THRESH_BINARY)
     cv2.imwrite("mask.png", mask)
-    
+
     # Generate satellite image
     satellite_image = generate_satellite_image("map_image.png", "mask.png", response['output']['feature_representation']['properties']['description'])
-    
+
     return map_image_path, satellite_image
 
 # Gradio interface
@@ -160,8 +167,8 @@ with gr.Blocks() as demo:
     with gr.Row():
         map_output = gr.Image(label="Map Visualization")
         satellite_output = gr.Image(label="Generated Satellite Image")
-    
+
     submit_btn.click(handle_query, inputs=[query_input], outputs=[map_output, satellite_output])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

requirements.txt

@@ -6,6 +6,7 @@ geopy       # For fetching geolocation data      # For PyTorch (used by Diffuser
 numpy  # For numerical operations
 diffusers
 transformers
+spaces
 torchvision
 opencv-python
 torch