add demo

app.py

@@ -2,17 +2,15 @@ import gradio as gr
 import torch
 from PIL import Image
 import requests
-from transformers import AutoProcessor, AutoModelForCausalLM
-from configuration_florence2 import Florence2Config
-from modeling_florence2 import Florence2ForCausalLM
-from processing_florence2 import Florence2Processor
+from transformers import AutoProcessor
+from modeling_florence2 import Florence2ForConditionalGeneration
 
 # Initialize model and processor
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-model = Florence2ForCausalLM.from_pretrained("PleIAs/Florence-PDF", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
-processor = Florence2Processor.from_pretrained("PleIAs/Florence-PDF", trust_remote_code=True)
+model = Florence2ForConditionalGeneration.from_pretrained("PleIAs/Florence-PDF", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
+processor = AutoProcessor.from_pretrained("PleIAs/Florence-PDF", trust_remote_code=True)
 
 # Define task prompts
 TASK_PROMPTS = {
@@ -31,8 +29,7 @@ def process_image(image, task):
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
     
     generated_ids = model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
+        **inputs,
         max_new_tokens=1024,
         num_beams=3,
         do_sample=False