Update app.py

app.py

@@ -6,12 +6,15 @@ import pickle
 import matplotlib.pyplot as plt
 import io
 from torch import nn
+from PIL import Image
 
 
+# Load ARIMA model
 with open("arima.pkl", "rb") as f:
     arima_model = pickle.load(f)
 
 
+# Define LSTM Model
 class LSTMModel(nn.Module):
     def __init__(self, input_size=1, hidden_size=50, num_layers=2, output_size=1):
         super(LSTMModel, self).__init__()
@@ -25,16 +28,20 @@ class LSTMModel(nn.Module):
         out = self.fc(out[:, -1, :])
         return out
 
+
 # Load trained LSTM
 lstm_model = LSTMModel()
 lstm_model.load_state_dict(torch.load("lstm.pth", map_location=torch.device('cpu')))
 lstm_model.eval()
 
 
+# ARIMA Prediction
 def predict_arima(values, horizon=10):
     forecast = arima_model.forecast(steps=horizon)
     return forecast.tolist()
 
+
+# LSTM Prediction
 def predict_lstm(values, horizon=10):
     seq = torch.tensor(values[-50:], dtype=torch.float32).view(1, -1, 1)
     preds = []
@@ -46,6 +53,7 @@ def predict_lstm(values, horizon=10):
     return preds
 
 
+# Forecast Function
 def forecast(file, horizon, model_choice):
     df = pd.read_csv(file.name)
     if "Close" not in df.columns:
@@ -78,13 +86,17 @@ def forecast(file, horizon, model_choice):
     plt.ylabel("Price")
     plt.legend()
 
+    # Save plot to buffer and convert to PIL
     buf = io.BytesIO()
     plt.savefig(buf, format="png")
     buf.seek(0)
-    
-    return forecast_df, buf
+    plt.close()
+    img = Image.open(buf)
+
+    return forecast_df, img
 
 
+# Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# 📈 Stock Price Forecasting Demo")
     gr.Markdown(