Update app.py

app.py

@@ -4,34 +4,26 @@ import re
 from PIL import Image 
 import os
 import numpy as np
-
 import spaces
 import subprocess
+import torch
+
+
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-model = AutoModelForCausalLM.from_pretrained('thwri/CogFlorence-2.1-Large', trust_remote_code=True).eval()#.to("cuda").eval()
-processor = AutoProcessor.from_pretrained('thwri/CogFlorence-2.1-Large', trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    'PJMixers-Dev/Florence-2-base-danbooru2022-316k',
+    trust_remote_code=True,
+    torch_dtype=torch.float32
+).eval()
+processor = AutoProcessor.from_pretrained(
+    'PJMixers-Dev/Florence-2-base-danbooru2022-316k',
+    trust_remote_code=True
+)
 
-TITLE = "# [thwri/CogFlorence-2.1-Large](https://huggingface.co/thwri/CogFlorence-2.1-Large/)"
-DESCRIPTION = "[microsoft/Florence-2-large](https://huggingface.co/microsoft/Florence-2-large) tuned on [Ejafa/ye-pop](https://huggingface.co/datasets/Ejafa/ye-pop) captioned with [CogVLM2](https://huggingface.co/THUDM/cogvlm2-llama3-chat-19B)"
+TITLE = "# [PJMixers-Dev/Florence-2-base-danbooru2022-316k](https://huggingface.co/PJMixers-Dev/Florence-2-base-danbooru2022-316k/)"
+DESCRIPTION = "[microsoft/Florence-2](https://huggingface.co/microsoft/Florence-2) tuned on [animelover/danbooru2022](https://huggingface.co/datasets/animelover/danbooru2022)."
 
-def modify_caption(caption: str) -> str:
-    special_patterns = [
-        (r'the image is ', ''),
-        (r'the image captures ', ''),
-        (r'the image showcases ', ''),
-        (r'the image shows ', ''),
-        (r'the image ', ''),
-    ]
-    
-    for pattern, replacement in special_patterns:
-        caption = re.sub(pattern, replacement, caption, flags=re.IGNORECASE)
-    
-    caption = caption.replace('\n', '').replace('\r', '')
-    caption = re.sub(r'(?<=[.,?!])(?=[^\s])', r' ', caption)
-    caption = ' '.join(caption.strip().splitlines())
-    
-    return caption
 
 @spaces.GPU
 def process_image(image):
@@ -42,19 +34,18 @@ def process_image(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
     
-    prompt = "<MORE_DETAILED_CAPTION>"
-    
-    inputs = processor(text=prompt, images=image, return_tensors="pt")#.to("cuda")
+    inputs = processor(text="<CAPTION>", images=image, return_tensors="pt").to(torch.float32)
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],
         pixel_values=inputs["pixel_values"],
         max_new_tokens=1024,
-        num_beams=3,
+        num_beams=5,
         do_sample=True
     )
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
-    return modify_caption(parsed_answer["<MORE_DETAILED_CAPTION>"])
+
+    return processor.post_process_generation(generated_text, task="<CAPTION>", image_size=(image.width, image.height))
+
 
 def extract_frames(image_path, output_folder):
     with Image.open(image_path) as img:
@@ -72,6 +63,7 @@ def extract_frames(image_path, output_folder):
         
         return frame_paths
 
+
 def process_folder(folder_path):
     if not os.path.isdir(folder_path):
         return "Invalid folder path."
@@ -119,6 +111,7 @@ def process_folder(folder_path):
                     processed_files.append(f"Processed {filename} -> {txt_filename}")
     
     result = "\n".join(processed_files + skipped_files)
+
     return result if result else "No image files found or all files were skipped in the specified folder."
 
 css = """