Update app.py

app.py

@@ -1,21 +1,23 @@
 import torch
-import gradio as gr
-from transformers import AutoProcessor, AutoModelForPreTraining
+import yfinance as yf
 import matplotlib.pyplot as plt
 import mplfinance as mpf
-import yfinance as yf
 from PIL import Image, ImageDraw, ImageFont
+import gradio as gr
 import datetime
-import tempfile
 import logging
+from transformers import AutoProcessor, AutoModelForPreTraining
+import tempfile
+import os
 
 # Configure logging
 logging.basicConfig(filename='debug.log', level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Load the model and processor
+# Load the ChartGemma model and processor
 processor = AutoProcessor.from_pretrained("mobenta/chart_analysis")
 model = AutoModelForPreTraining.from_pretrained("mobenta/chart_analysis")
 
+@spaces.GPU
 def predict(image, input_text):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
@@ -23,7 +25,7 @@ def predict(image, input_text):
     image = image.convert("RGB")
     inputs = processor(text=input_text, images=image, return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
-    
+
     prompt_length = inputs['input_ids'].shape[1]
     generate_ids = model.generate(**inputs, max_new_tokens=512)
     output_text = processor.batch_decode(generate_ids[:, prompt_length:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
@@ -50,10 +52,10 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d', indic
     try:
         logging.debug(f"Creating chart for {ticker} with timeframe {timeframe} and saving to {filename}")
         title = f"{ticker.upper()} Price Data (Timeframe: {timeframe})"
-        
+
         plt.rcParams["axes.titlesize"] = 10
         my_style = mpf.make_mpf_style(base_mpf_style='charles')
-        
+
         # Calculate indicators if selected
         addplot = []
         if indicators:
@@ -93,7 +95,7 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d', indic
 
         fig, axlist = mpf.plot(data, type='candle', style=my_style, volume=True, addplot=addplot, returnfig=True)
         fig.suptitle(title, y=0.98)
-        
+
         # Save chart image
         fig.savefig(filename, dpi=300)
         plt.close(fig)
@@ -117,7 +119,7 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d', indic
             metrics["SMA 50"] = f"${data['Close'].rolling(window=50).mean().iloc[-1]:,.2f}"
         if 'SMA200' in indicators:
             metrics["SMA 200"] = f"${data['Close'].rolling(window=200).mean().iloc[-1]:,.2f}"
-        
+
         # Draw metrics on the image
         y_text = image.height - 50  # Starting y position for text
         for key, value in metrics.items():
@@ -130,23 +132,47 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d', indic
         resized_image = image.resize(new_size, Image.LANCZOS)
         resized_image.save(filename)
 
-        logging.debug(f"Resized image saved to {filename}")
-        return filename
+        logging.debug(f"Resized image with timeframe {timeframe} and ticker {ticker} saved to {filename}")
+    except Exception as e:
+        logging.error(f"Error creating or resizing chart: {e}")
+        raise
+
+def combine_images(image_paths, output_path='combined_chart.png'):
+    try:
+        logging.debug(f"Combining images {image_paths} into {output_path}")
+        images = [Image.open(path) for path in image_paths]
+
+        # Calculate total width and max height for combined image
+        total_width = sum(img.width for img in images)
+        max_height = max(img.height for img in images)
+
+        combined_image = Image.new('RGB', (total_width, max_height))
+        x_offset = 0
+        for img in images:
+            combined_image.paste(img, (x_offset, 0))
+            x_offset += img.width
+
+        combined_image.save(output_path)
+        logging.debug(f"Combined image saved to {output_path}")
+        return output_path
     except Exception as e:
-        logging.error(f"Error creating stock chart: {e}")
+        logging.error(f"Error combining images: {e}")
         raise
 
 def gradio_interface(ticker1, ticker2, ticker3, ticker4, start_date, end_date, query, analysis_type, interval, indicators):
     try:
-        chart_paths = []
+        logging.debug(f"Starting gradio_interface with tickers: {ticker1}, {ticker2}, {ticker3}, {ticker4}, start_date: {start_date}, end_date: {end_date}, query: {query}, analysis_type: {analysis_type}, interval: {interval}")
+
         tickers = [ticker1, ticker2, ticker3, ticker4]
+        chart_paths = []
 
-        for ticker in tickers:
-            if ticker.strip():
-                data = fetch_stock_data(ticker, start_date, end_date, interval)
-                chart_path = f"{ticker}_chart.png"
-                create_stock_chart(data, ticker, filename=chart_path, timeframe=interval, indicators=indicators)
-                chart_paths.append(chart_path)
+        for i, ticker in enumerate(tickers):
+            if ticker:
+                data = fetch_stock_data(ticker, start=start_date, end=end_date, interval=interval)
+                with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_chart:
+                    chart_path = temp_chart.name
+                    create_stock_chart(data, ticker, chart_path, timeframe=interval, indicators=indicators)
+                    chart_paths.append(chart_path)
 
         if analysis_type == 'Comparative Analysis' and len(chart_paths) > 1:
             with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_combined_chart:
@@ -155,6 +181,7 @@ def gradio_interface(ticker1, ticker2, ticker3, ticker4, start_date, end_date, q
                 insights = predict(Image.open(combined_chart_path), query)
                 return insights, combined_chart_path
 
+        # No comparative analysis, just return the single chart
         if chart_paths:
             insights = predict(Image.open(chart_paths[0]), query)
             return insights, chart_paths[0]