Update app.py

app.py

@@ -6,23 +6,40 @@ from PIL import Image
 import gradio as gr
 import logging
 from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+import spaces
 
 # Configure logging to write to a file
 logging.basicConfig(filename='debug.log', level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Check GPU availability and initialize device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print("CUDA Available:", torch.cuda.is_available())
-print("Using device:", device)
+# Load the ChartGemma model and processor outside the GPU context
+model = PaliGemmaForConditionalGeneration.from_pretrained("ahmed-masry/chartgemma")
+processor = AutoProcessor.from_pretrained("ahmed-masry/chartgemma")
 
-# Load the ChartGemma model and processor
-try:
-    model = PaliGemmaForConditionalGeneration.from_pretrained("ahmed-masry/chartgemma", torch_dtype=torch.float16).to(device)
-    processor = AutoProcessor.from_pretrained("ahmed-masry/chartgemma")
-    print("Model and processor loaded successfully.")
-except Exception as e:
-    print("Error loading model or processor:", e)
-    logging.error(f"Error loading model or processor: {e}")
+@spaces.GPU
+def predict(image, input_text, ticker1=None, ticker2=None):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
+    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 insights using the model
+    generate_ids = model.generate(**inputs, num_beams=4, max_new_tokens=512)
+    output_text = processor.batch_decode(generate_ids[:, prompt_length:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+
+    # Replace placeholders with actual ticker names in the insights
+    if ticker1:
+        output_text = output_text.replace("[First Ticker]", ticker1)
+    if ticker2:
+        output_text = output_text.replace("[Second Ticker]", ticker2)
+
+    logging.debug(f"Generated insights: {output_text}")
+
+    return output_text
 
 # Function to fetch stock data with different intervals
 def fetch_stock_data(ticker='TSLA', start='2023-01-01', end='2024-01-01', interval='1d'):
@@ -80,36 +97,6 @@ def combine_images(image1_path, image2_path, output_path='combined_chart.png'):
         logging.error(f"Error combining images: {e}")
         raise
 
-# Function to generate insights
-def generate_insights(image, query, ticker1=None, ticker2=None):
-    try:
-        logging.debug(f"Generating insights for query: {query}")
-
-        # Open and process the image
-        image = Image.open(image).convert('RGB')
-        inputs = processor(text=query, images=image, return_tensors="pt")
-        logging.debug(f"Inputs prepared with shapes {inputs['input_ids'].shape} and {inputs['pixel_values'].shape}")
-        
-        prompt_length = inputs['input_ids'].shape[1]
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-
-        # Generate insights using the model
-        generate_ids = model.generate(**inputs, num_beams=4, max_new_tokens=512)
-        output_text = processor.batch_decode(generate_ids[:, prompt_length:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-
-        # Replace placeholders with actual ticker names in the insights
-        if ticker1:
-            output_text = output_text.replace("[First Ticker]", ticker1)
-        if ticker2:
-            output_text = output_text.replace("[Second Ticker]", ticker2)
-
-        logging.debug(f"Generated insights: {output_text}")
-
-        return output_text
-    except Exception as e:
-        logging.error(f"Error generating insights: {e}")
-        return f"Error generating insights: {e}"
-
 # Function to handle the Gradio interface
 def gradio_interface(ticker1, start_date, end_date, ticker2, query, analysis_type, interval):
     try:
@@ -127,10 +114,10 @@ def gradio_interface(ticker1, start_date, end_date, ticker2, query, analysis_typ
 
             # Combine the two charts into one image
             combined_chart_path = combine_images(chart_path1, chart_path2)
-            insights = generate_insights(combined_chart_path, query, ticker1, ticker2)
+            insights = predict(Image.open(combined_chart_path), query, ticker1, ticker2)
             return insights, combined_chart_path
 
-        insights = generate_insights(chart_path1, query, ticker1)
+        insights = predict(Image.open(chart_path1), query, ticker1)
         return insights, chart_path1
     except Exception as e:
         logging.error(f"Error processing image or query: {e}")