Update app.py

app.py

@@ -14,48 +14,15 @@ import os
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-model_dir = 'medieval-data/florence2-medieval-bbox-line-detection'
 model_dir = "medieval-data/florence2-medieval-bbox-zone-detection"
 
-def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
-    """Work around for https://huggingface.co/microsoft/phi-1_5/discussions/72."""
-    if not str(filename).endswith("/modeling_florence2.py"):
-        return get_imports(filename)
-    imports = get_imports(filename)
-    imports.remove("flash_attn")
-    return imports
-    
-with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
-    # Load the configuration
-    config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
-    
-    # Modify the vision configuration
-    if hasattr(config, 'vision_config'):
-        config.vision_config.model_type = 'davit'
-    
-    print("Modified vision configuration:")
-    print(config.vision_config)
-    
-    # Try to load the model with the modified configuration
-    try:
-        model = AutoModelForCausalLM.from_pretrained(
-            model_dir,
-            config=config,
-            trust_remote_code=True
-        )
-        print("Model loaded successfully!")
-    except Exception as e:
-        print(f"Failed to load model: {str(e)}")
-
-    # Load the processor without specifying a revision
-    try:
-        processor = AutoProcessor.from_pretrained(
-            model_dir, 
-            trust_remote_code=True
-        )
-        print("Processor loaded successfully!")
-    except Exception as e:
-        print(f"Failed to load processor: {str(e)}")
+# Load the configuration
+config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+        model_dir,
+        trust_remote_code=True
+    )
+
 TITLE = "# [Florence-2- Medieval Manuscript Layout Parsing Demo](https://huggingface.co/medieval-data/florence2-medieval-bbox-zone-detection)"
 DESCRIPTION = "The demo for Florence-2 fine-tuned on CATMuS Segmentation Dataset. This app has two models: one for line detection and one for zone detection."
 
@@ -63,7 +30,7 @@ DESCRIPTION = "The demo for Florence-2 fine-tuned on CATMuS Segmentation Dataset
 colormap = plt.cm.get_cmap('tab20')
 
 @spaces.GPU
-def process_image(image, text_input=None):
+def process_image(image):
     max_size = 1000
     prompt = "<OD>"
 
@@ -111,14 +78,14 @@ def visualize_bboxes(result, image):
     plt.axis('off')
     return fig
 
-def run_example(image, text_input=None):
+def run_example(image):
     if isinstance(image, str):  # If image is a URL
         response = requests.get(image)
         image = Image.open(BytesIO(response.content))
     elif isinstance(image, np.ndarray):  # If image is a numpy array
         image = Image.fromarray(image)
     
-    result, processed_image = process_image(image, text_input)
+    result, processed_image = process_image(image)
     fig = visualize_bboxes(result, processed_image)
     
     # Convert matplotlib figure to image
@@ -128,15 +95,13 @@ def run_example(image, text_input=None):
     output_image = Image.open(img_buf)
     
     return output_image
-
 css = """
   #output {
-    height: 500px; 
+    height: 1000px; 
     overflow: auto; 
     border: 1px solid #ccc; 
   }
 """
-
 with gr.Blocks(css=css) as demo:
     gr.Markdown(TITLE)
     gr.Markdown(DESCRIPTION)
@@ -144,23 +109,19 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             with gr.Column():
                 input_img = gr.Image(label="Input Picture")
-                text_input = gr.Textbox(label="Text Input (optional)")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():
                 output_img = gr.Image(label="Output Image with Bounding Boxes")
         gr.Examples(
             examples=[
-                ["https://huggingface.co/datasets/CATMuS/medieval-segmentation/resolve/main/data/dev/london-british-library-egerton-821/page-002-of-004.jpg", None],
-                ["https://huggingface.co/datasets/CATMuS/medieval-segmentation/resolve/main/data/dev/paris-bnf-lat-12449/page-002-of-003.jpg", None],
-                ["https://huggingface.co/datasets/CATMuS/medieval-segmentation/resolve/main/data/dev/paris-bnf-nal-1909/page-009-of-012.jpg", None],
-                ["https://huggingface.co/datasets/CATMuS/medieval-segmentation/resolve/main/data/test/paris-bnf-fr-574/page-001-of-003.jpg", None]
+                ["https://huggingface.co/datasets/CATMuS/medieval-segmentation/resolve/main/data/dev/london-british-library-egerton-821/page-002-of-004.jpg"],
             ],
-            inputs=[input_img, text_input],
+            inputs=[input_img],
             outputs=[output_img],
             fn=run_example,
             cache_examples=True,
             label='Try the examples below'
         )
-        submit_btn.click(run_example, [input_img, text_input], [output_img])
+        submit_btn.click(run_example, [input_img], [output_img])
 
 demo.launch(debug=True)