Update app.py

app.py

@@ -6,17 +6,18 @@ import uvicorn
 import re
 from dotenv import load_dotenv
 from spaces import GPU
+from spaces.zero import ZeroGPU
 
 load_dotenv()
 
-app = FastAPI()
-
 # Initialize ZeroGPU
 try:
-    GPU.initialize()
+    ZeroGPU.initialize()
 except Exception as e:
     print(f"ZeroGPU initialization failed: {e}")
 
+app = FastAPI()
+
 # Global data dictionary
 global_data = {
     'models': {},
@@ -120,16 +121,22 @@ def remove_repetitive_responses(responses):
             unique_responses.append({'model': response['model'], 'response': normalized_response})
     return unique_responses
 
-@app.post("/generate/")
-@GPU(duration=0)
+def generate_model_response(model, inputs, top_k, top_p, temperature):
+    try:
+        return model.generate(inputs, top_k=top_k, top_p=top_p, temperature=temperature)
+    except Exception as e:
+        print(f"Error generating response with model {model}: {e}")
+        return ""
+
+@app.post("/generate")
 async def generate(request: ChatRequest):
     try:
         inputs = normalize_input(request.message)
         futures = [
-            executor.submit(model.generate, inputs, top_k=request.top_k, top_p=request.top_p, temperature=request.temperature)
+            executor.submit(generate_model_response, model, inputs, request.top_k, request.top_p, request.temperature)
             for model in global_data['models'].values()
         ]
-        responses = [{'model': model, 'response': future.result()} for model, future in zip(global_data['models'].keys(), as_completed(futures))]
+        responses = [{'model': model_name, 'response': future.result()} for model_name, future in zip(global_data['models'].keys(), as_completed(futures))]
         unique_responses = remove_repetitive_responses(responses)
         return unique_responses
     except Exception as e: