trying to use state to save the output

app.py

@@ -149,7 +149,7 @@ def run_example(image, model_id="Qwen/Qwen2-VL-7B-Instruct", run_ner=False, ner_
         
         # Create TextWithMetadata instance with the highlighted text and metadata
         result = TextWithMetadata(highlighted_text, original_text=ocr_text, entities=ner_results)
-        return result
+        return result, result  # Return twice: once for display, once for state
     
     # If NER is disabled, return the text without highlighting
     result = TextWithMetadata([(ocr_text, None)], original_text=ocr_text, entities=[])
@@ -227,6 +227,9 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
+    # Add state variables to store OCR results
+    ocr_state = gr.State()
+    
     gr.Image("Caracal.jpg", interactive=False)
     with gr.Tab(label="Image Input", elem_classes="tabs"):
         with gr.Row():
@@ -254,20 +257,21 @@ with gr.Blocks(css=css) as demo:
             outputs=[ner_labels]
         )
         
-        # Update submit button click handler
+        # Modify the submit button click handler to update state
         submit_btn.click(
             run_example,
             inputs=[input_img, model_selector, ner_checkbox, ner_labels],
-            outputs=[output_text]
+            outputs=[output_text, ocr_state]  # Add ocr_state to outputs
         )
     with gr.Row():
         filename = gr.Textbox(label="Save filename (without extension)", placeholder="Enter filename to save")
         download_btn = gr.Button("Download Image & Text", elem_classes="submit-btn")
         download_output = gr.File(label="Download")
 
-    def create_zip(image, text_data, fname):
+    # Modify create_zip to use the state data
+    def create_zip(image, text_data, fname, ocr_result):
         # Validate inputs
-        if not fname or not text_data or not isinstance(image, (Image.Image, np.ndarray)):
+        if not fname or not image or not isinstance(image, (Image.Image, np.ndarray)):
             return None
         
         try:
@@ -280,9 +284,9 @@ with gr.Blocks(css=css) as demo:
                 img_path = os.path.join(temp_dir, f"{fname}.png")
                 image.save(img_path)
                 
-                # Get the original OCR text and entities
-                original_text = getattr(text_data, 'original_text', '')
-                entities = getattr(text_data, 'entities', [])
+                # Use the OCR result from state
+                original_text = ocr_result.original_text if ocr_result else ""
+                entities = ocr_result.entities if ocr_result else []
                 
                 # Save text
                 txt_path = os.path.join(temp_dir, f"{fname}.txt")
@@ -317,10 +321,10 @@ with gr.Blocks(css=css) as demo:
             print(f"Error creating zip: {str(e)}")
             return None
 
-    # Add the click handler for the download button
+    # Update the download button click handler to include state
     download_btn.click(
         create_zip,
-        inputs=[input_img, output_text, filename],
+        inputs=[input_img, output_text, filename, ocr_state],
         outputs=[download_output]
     )