Update app.py

app.py

@@ -1,30 +1,34 @@
+import os
+os.environ['CUDA_HOME'] = '/usr/local/cuda'
+os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
+
 import gradio as gr
-import json
 import torch
+import json
 from huggingface_hub import InferenceClient
-from huggingface_hub import spaces
+import spaces
 
-# Initialize the InferenceClient for GPU and CPU
-client_gpu = InferenceClient(model="gpt-3.5-turbo")  # Assuming this will use GPU by default if available
-client_cpu = InferenceClient(model="gpt-3.5-turbo")  # Duplicate client for CPU fallback
+# Initialize the InferenceClient
+client = InferenceClient(model="gpt-3.5-turbo")
 
-# Function to attempt to fix JSON using the appropriate resource
+@spaces.GPU(duration=120)
 def ai_fix_json(json_data):
     prompt = f"Fix the following JSON data and make it valid:\n\n{json_data}\n\nFixed JSON:"
     
     try:
-        # Check if the space is running on GPU
         if torch.cuda.is_available():
-            response = client_gpu.text_generation(prompt, max_new_tokens=1024)
-            return response[0]['generated_text'].split("Fixed JSON:")[-1].strip(), "JSON fixed using AI on GPU."
+            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."
         else:
             raise RuntimeError("GPU not available, falling back to CPU.")
     
     except Exception as gpu_error:
         # Fall back to CPU if GPU processing fails
         print(f"Falling back to CPU due to: {gpu_error}")
-        response = client_cpu.text_generation(prompt, max_new_tokens=1024)
-        return response[0]['generated_text'].split("Fixed JSON:")[-1].strip(), "JSON fixed using AI on CPU."
+        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."
 
 def process_file(uploaded_file):
     json_data = uploaded_file  # This is already the content of the file as a string
@@ -41,8 +45,9 @@ iface = gr.Interface(
     fn=process_file,
     inputs=gr.File(label="Upload your JSON file"),
     outputs=[gr.JSON(label="Fixed JSON"), "text", gr.File(label="Download cleaned JSON file")],
-    title="AI-Powered JSON Cleaner with Dynamic Resource Allocation",
+    title="AI-Powered JSON Cleaner with GPU Support",
     description="Upload a JSON file to automatically fix, remove duplicates, and download the cleaned version using AI with GPU/CPU fallback."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()