Update app.py

app.py

@@ -2,13 +2,8 @@ from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 import gradio as gr
 from PIL import Image
 
-# Use a public model identifier; change this if you have a different one or want to use a private model.
+# Use a public model identifier. If you need a private model, remember to authenticate.
 model_name = "google/pix2struct-textcaps-base"
-
-# If you need authentication for a private repo, pass the token as follows:
-# model = Pix2StructForConditionalGeneration.from_pretrained(model_name, use_auth_token="YOUR_TOKEN")
-# processor = Pix2StructProcessor.from_pretrained(model_name, use_auth_token="YOUR_TOKEN")
-
 model = Pix2StructForConditionalGeneration.from_pretrained(model_name)
 processor = Pix2StructProcessor.from_pretrained(model_name)
 
@@ -17,8 +12,7 @@ def solve_math_problem(image):
         # Ensure the image is in RGB format.
         image = image.convert("RGB")
         
-        # Preprocess the image and text.
-        # Note: We omit header_text since this is not a VQA task.
+        # Preprocess the image and text. Note that header_text is omitted as it's not used for non-VQA tasks.
         inputs = processor(
             images=[image],
             text="Solve the following math problem:",
@@ -26,7 +20,7 @@ def solve_math_problem(image):
             max_patches=2048
         )
         
-        # Generate the solution with specified generation parameters.
+        # Generate the solution with generation parameters.
         predictions = model.generate(
             **inputs,
             max_new_tokens=200,
@@ -35,7 +29,7 @@ def solve_math_problem(image):
             temperature=0.2
         )
         
-        # Decode the problem text and the generated solution.
+        # Decode the problem text and generated solution.
         problem_text = processor.decode(
             inputs["input_ids"][0],
             skip_special_tokens=True,
@@ -58,8 +52,7 @@ demo = gr.Interface(
     inputs=gr.Image(
         type="pil", 
         label="Upload Handwritten Math Problem",
-        image_mode="RGB",
-        source="upload"
+        image_mode="RGB"  # This forces the input to be RGB.
     ),
     outputs=gr.Textbox(label="Solution", show_copy_button=True),
     title="Handwritten Math Problem Solver",