Update app.py

app.py

@@ -1,13 +1,20 @@
 import os
-os.environ['CUDA_HOME'] = '/usr/local/cuda'
-os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
-
 import gradio as gr
 import torch
 import json
 from huggingface_hub import InferenceClient
 import spaces
 
+# Ensure CUDA is configured correctly
+os.environ['CUDA_HOME'] = '/usr/local/cuda'
+os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
+
+def initialize_zerogpu():
+    # Example function to initialize ZeroGPU, replace with actual initialization if needed
+    pass
+
+# Function to perform JSON correction using a chosen model
+@spaces.GPU(duration=120)  # Use GPU for this function, if available
 def ai_fix_json(model_id, json_data):
     # Initialize the InferenceClient with the chosen model
     client = InferenceClient(model=model_id)
@@ -15,10 +22,11 @@ def ai_fix_json(model_id, json_data):
     prompt = f"Fix the following JSON data and make it valid:\n\n{json_data}\n\nFixed JSON:"
     
     try:
+        # Check if GPU is available
         if torch.cuda.is_available():
             response = client.text_generation(prompt, max_new_tokens=1024)
             fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
-            return fixed_json, "JSON fixed using AI on GPU with model: " + model_id
+            return fixed_json, f"JSON fixed using AI on GPU with model: {model_id}"
         else:
             raise RuntimeError("GPU not available, falling back to CPU.")
     
@@ -27,7 +35,7 @@ def ai_fix_json(model_id, json_data):
         print(f"Falling back to CPU due to: {gpu_error}")
         response = client.text_generation(prompt, max_new_tokens=1024)
         fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
-        return fixed_json, "JSON fixed using AI on CPU with model: " + model_id
+        return fixed_json, f"JSON fixed using AI on CPU with model: {model_id}"
 
 def process_file(model_id, uploaded_file):
     json_data = uploaded_file  # This is already the content of the file as a string
@@ -52,4 +60,10 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
+    try:
+        initialize_zerogpu()
+        print("ZeroGPU initialized.")
+    except Exception as e:
+        print(f"ZeroGPU initialization failed: {e}. Falling back to CPU.")
+    
     iface.launch()