Update app.py

app.py

@@ -2,8 +2,13 @@ from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 import gradio as gr
 from PIL import Image
 
-# Load the pre-trained Pix2Struct model and processor.
-model_name = "google/pix2struct-mathqa-base"
+# Use a public model identifier; change this if you have a different one or want to use a private model.
+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)
 
@@ -13,12 +18,12 @@ def solve_math_problem(image):
         image = image.convert("RGB")
         
         # Preprocess the image and text.
-        # Note: We omit the header_text parameter because this is not a VQA task.
+        # Note: We omit header_text since this is not a VQA task.
         inputs = processor(
-            images=[image],                           # Provide a list of images.
-            text="Solve the following math problem:", # Prompt text.
+            images=[image],
+            text="Solve the following math problem:",
             return_tensors="pt",
-            max_patches=2048                          # Increase the maximum patches for better math handling.
+            max_patches=2048
         )
         
         # Generate the solution with specified generation parameters.
@@ -30,8 +35,7 @@ def solve_math_problem(image):
             temperature=0.2
         )
         
-        # Decode the input text and the model prediction.
-        # Here, we access "input_ids" via the dictionary key.
+        # Decode the problem text and the generated solution.
         problem_text = processor.decode(
             inputs["input_ids"][0],
             skip_special_tokens=True,
@@ -54,14 +58,14 @@ demo = gr.Interface(
     inputs=gr.Image(
         type="pil", 
         label="Upload Handwritten Math Problem",
-        image_mode="RGB",  # Force RGB conversion.
+        image_mode="RGB",
         source="upload"
     ),
     outputs=gr.Textbox(label="Solution", show_copy_button=True),
     title="Handwritten Math Problem Solver",
     description="Upload an image of a handwritten math problem (algebra, arithmetic, etc.) and get the solution",
     examples=[
-        ["example_addition.png"],  # Ensure these example files exist in your working directory.
+        ["example_addition.png"],
         ["example_algebra.jpg"]
     ],
     theme="soft",