Update app.py

app.py

@@ -12,14 +12,24 @@ import subprocess
 # Install necessary packages
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-# Initialize Florence model
-model_id = 'microsoft/Florence-2-large'
-florence_model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda").eval()
-florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-
 # Initialize Llama Cleaner model
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
+# Define available models
+available_models = [
+    'microsoft/Florence-2-base',
+    'microsoft/Florence-2-base-ft',
+    'microsoft/Florence-2-large',
+    'microsoft/Florence-2-large-ft'
+]
+
+# Load all models and processors
+model_dict = {}
+for model_id in available_models:
+    florence_model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda").eval()
+    florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+    model_dict[model_id] = (florence_model, florence_processor)
+
 @spaces.GPU()
 def process_image(image, mask, strategy, sampler, fx=1, fy=1):
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
@@ -45,7 +55,7 @@ def process_image(image, mask, strategy, sampler, fx=1, fy=1):
 def create_mask(image, prediction):
     mask = Image.new("RGBA", image.size, (0, 0, 0, 255))  # Black background
     draw = ImageDraw.Draw(mask)
-    scale = 1
+    scale = 1.1
     for polygons in prediction['polygons']:
         for _polygon in polygons:
             _polygon = np.array(_polygon).reshape(-1, 2)
@@ -56,12 +66,14 @@ def create_mask(image, prediction):
     return mask
 
 @spaces.GPU()
-def process_images_florence_lama(image):
+def process_images_florence_lama(image, model_choice):
+    florence_model, florence_processor = model_dict[model_choice]
+    
     # Convert image to OpenCV format
     image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
     
     # Run Florence to get mask
-    text_input = 'watermark'  # Teks untuk Florence agar mengenali watermark
+    text_input = 'watermark, text'  # Teks untuk Florence agar mengenali watermark
     task_prompt = '<REGION_TO_SEGMENTATION>'
     image_pil = Image.fromarray(image_cv)  # Convert array to PIL Image
     inputs = florence_processor(text=task_prompt + text_input, images=image_pil, return_tensors="pt").to("cuda")
@@ -92,11 +104,14 @@ def process_images_florence_lama(image):
 # Define Gradio interface
 demo = gr.Interface(
     fn=process_images_florence_lama,
-    inputs=gr.Image(type="pil", label="Input Image"),
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Dropdown(choices=available_models, value='microsoft/Florence-2-large', label="Choose Florence Model")
+    ],
     outputs=gr.Image(type="pil", label="Output Image"),
-    title="Watermark Remover.",
+    title="Watermark Remover",
     description="Upload images and remove selected watermarks using Florence and Lama Cleaner.\nhttps://github.com/Damarcreative/rem-wm.git"
 )
-# Launch Gradio interface with example images
+
 if __name__ == "__main__":
     demo.launch()