Update app.py

app.py

@@ -1,38 +1,33 @@
-import os
 import torch
-from flask import Flask, request, jsonify
+import gradio as gr
 import numpy as np
-from src.model import LSTM
-
-# Initialize Flask app
-app = Flask(__name__)
-
-# Device setup
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-# Ensure the model file is available in the environment
-model_path = './water_forecast_2.pth'
-if not os.path.exists(model_path):
-    raise FileNotFoundError(f"Model file '{model_path}' not found.")
+from src.model import LSTM  # Adjust to your model path
 
 # Load the model
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model_path = "./water_forecast_2.pth"  # Your model path
 model = LSTM(input_size=8, lstm_layer_sizes=[128, 128, 128], output_size=3).to(device)
-
-print("Loading model...")
 model.load_state_dict(torch.load(model_path, map_location=device))
-print("Model loaded successfully")
 model.eval()
 
-@app.route("/predict", methods=["POST"])
-def predict_usage():
-    data = request.get_json()
-    structured_data = data["structured_data"]
+# Define the prediction function
+def predict_water_usage(state_idx, target_year, structured_data):
+    # Convert input data to tensor
     tensor_data = torch.tensor(np.array(list(structured_data.values())), dtype=torch.float32).to(device)
-
     with torch.no_grad():
-        outputs = model(tensor_data)
+        output = model(tensor_data)
+    return output.tolist()
+
+# Set up Gradio interface
+inputs = [
+    gr.inputs.Number(label="State Index"),
+    gr.inputs.Number(label="Target Year"),
+    gr.inputs.JSON(label="Structured Data")  # Expects JSON input
+]
+outputs = gr.outputs.JSON(label="Prediction")
 
-    return jsonify({"prediction": outputs.tolist()})
+interface = gr.Interface(fn=predict_water_usage, inputs=inputs, outputs=outputs)
 
+# Launch the Gradio app
 if __name__ == "__main__":
-    app.run(debug=True)
+    interface.launch()