Florence sam flux first commit

.gitignore

@@ -0,0 +1,3 @@
+/venv
+/.idea
+/tmp

README.md

@@ -0,0 +1,12 @@
+---
+title: Florence2 + SAM2 + FLUX
+emoji: 🔥
+colorFrom: purple
+colorTo: green
+sdk: gradio
+sdk_version: 4.40.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,121 @@
+import torch
+import numpy as np
+from PIL import Image
+from diffusers import FluxInpaintPipeline
+from utils.florence import load_florence_model, run_florence_inference, FLORENCE_OPEN_VOCABULARY_DETECTION_TASK
+from utils.sam import load_sam_image_model, run_sam_inference
+import gradio as gr
+import supervision as sv
+
+# Load models
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+FLUX_PIPE = FluxInpaintPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(
+    DEVICE)
+FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE)
+SAM_MODEL = load_sam_image_model(device=DEVICE)
+
+COLORS = ['#FF1493', '#00BFFF', '#FF6347', '#FFD700', '#32CD32', '#8A2BE2']
+COLOR_PALETTE = sv.ColorPalette.from_hex(COLORS)
+BOX_ANNOTATOR = sv.BoxAnnotator(color=COLOR_PALETTE, color_lookup=sv.ColorLookup.INDEX)
+LABEL_ANNOTATOR = sv.LabelAnnotator(
+    color=COLOR_PALETTE,
+    color_lookup=sv.ColorLookup.INDEX,
+    text_position=sv.Position.CENTER_OF_MASS,
+    text_color=sv.Color.from_hex("#000000"),
+    border_radius=5
+)
+MASK_ANNOTATOR = sv.MaskAnnotator(
+    color=COLOR_PALETTE,
+    color_lookup=sv.ColorLookup.INDEX
+)
+
+
+def visualize_detections(image, detections):
+    output_image = image.copy()
+    output_image = MASK_ANNOTATOR.annotate(output_image, detections)
+    output_image = BOX_ANNOTATOR.annotate(output_image, detections)
+    output_image = LABEL_ANNOTATOR.annotate(output_image, detections)
+    return output_image
+
+
+def detect_objects(image, text_prompt):
+    # Use Florence for object detection
+    _, result = run_florence_inference(
+        model=FLORENCE_MODEL,
+        processor=FLORENCE_PROCESSOR,
+        device=DEVICE,
+        image=image,
+        task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK,
+        text=text_prompt
+    )
+    detections = sv.Detections.from_lmm(
+        lmm=sv.LMM.FLORENCE_2,
+        result=result,
+        resolution_wh=image.size
+    )
+
+    # Use SAM to refine masks
+    detections = run_sam_inference(SAM_MODEL, image, detections)
+    return detections
+
+
+def inpaint_selected_objects(image, detections, selected_indices, inpaint_prompt):
+    mask = np.zeros(image.size[::-1], dtype=np.uint8)
+    for idx in selected_indices:
+        mask |= detections.mask[idx]
+
+    mask_image = Image.fromarray(mask * 255)
+
+    result = FLUX_PIPE(
+        prompt=inpaint_prompt,
+        image=image,
+        mask_image=mask_image,
+        num_inference_steps=30,
+        strength=0.85,
+    ).images[0]
+
+    return result
+
+
+def process_image(input_image, detection_prompt, inpaint_prompt, selected_objects):
+    detections = detect_objects(input_image, detection_prompt)
+
+    # Visualize detected objects
+    detected_image = visualize_detections(input_image, detections)
+
+    if selected_objects:
+        selected_indices = [int(idx) for idx in selected_objects.split(',')]
+        inpainted_image = inpaint_selected_objects(input_image, detections, selected_indices, inpaint_prompt)
+        return detected_image, inpainted_image
+    else:
+        return detected_image, None
+
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Object Detection and Inpainting with FLUX, Florence, and SAM")
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            detection_prompt = gr.Textbox(label="Detection Prompt", placeholder="Enter objects to detect")
+            detect_button = gr.Button("Detect Objects")
+        with gr.Column():
+            detected_image = gr.Image(type="pil", label="Detected Objects")
+            selected_objects = gr.Textbox(label="Selected Objects",
+                                          placeholder="Enter indices of objects to inpaint (comma-separated)")
+            inpaint_prompt = gr.Textbox(label="Inpainting Prompt", placeholder="Describe what to inpaint")
+            inpaint_button = gr.Button("Inpaint Selected Objects")
+    output_image = gr.Image(type="pil", label="Inpainted Result")
+
+    detect_button.click(
+        fn=lambda img, prompt: process_image(img, prompt, "", "")[0],
+        inputs=[input_image, detection_prompt],
+        outputs=detected_image
+    )
+    inpaint_button.click(
+        fn=process_image,
+        inputs=[input_image, detection_prompt, inpaint_prompt, selected_objects],
+        outputs=[detected_image, output_image]
+    )
+
+demo.launch(debug=False, show_error=True)

requirements.txt

@@ -0,0 +1,13 @@
+tqdm
+einops
+spaces
+timm
+transformers
+samv2
+gradio
+supervision
+opencv-python
+pytest
+torch
+numpy
+diffusers

utils/florence.py

@@ -0,0 +1,58 @@
+import os
+from typing import Union, Any, Tuple, Dict
+from unittest.mock import patch
+
+import torch
+from PIL import Image
+from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers.dynamic_module_utils import get_imports
+
+FLORENCE_CHECKPOINT = "microsoft/Florence-2-base"
+FLORENCE_OBJECT_DETECTION_TASK = '<OD>'
+FLORENCE_DETAILED_CAPTION_TASK = '<MORE_DETAILED_CAPTION>'
+FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK = '<CAPTION_TO_PHRASE_GROUNDING>'
+FLORENCE_OPEN_VOCABULARY_DETECTION_TASK = '<OPEN_VOCABULARY_DETECTION>'
+FLORENCE_DENSE_REGION_CAPTION_TASK = '<DENSE_REGION_CAPTION>'
+
+
+def fixed_get_imports(filename: Union[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
+
+
+def load_florence_model(
+    device: torch.device, checkpoint: str = FLORENCE_CHECKPOINT
+) -> Tuple[Any, Any]:
+    with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+        model = AutoModelForCausalLM.from_pretrained(
+            checkpoint, trust_remote_code=True).to(device).eval()
+        processor = AutoProcessor.from_pretrained(
+            checkpoint, trust_remote_code=True)
+        return model, processor
+
+
+def run_florence_inference(
+    model: Any,
+    processor: Any,
+    device: torch.device,
+    image: Image,
+    task: str,
+    text: str = ""
+) -> Tuple[str, Dict]:
+    prompt = task + text
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        num_beams=3
+    )
+    generated_text = processor.batch_decode(
+        generated_ids, skip_special_tokens=False)[0]
+    response = processor.post_process_generation(
+        generated_text, task=task, image_size=image.size)
+    return generated_text, response

utils/sam.py

@@ -0,0 +1,45 @@
+from typing import Any
+
+import numpy as np
+import supervision as sv
+import torch
+from PIL import Image
+from sam2.build_sam import build_sam2, build_sam2_video_predictor
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+SAM_CHECKPOINT = "checkpoints/sam2_hiera_small.pt"
+SAM_CONFIG = "sam2_hiera_s.yaml"
+
+
+def load_sam_image_model(
+    device: torch.device,
+    config: str = SAM_CONFIG,
+    checkpoint: str = SAM_CHECKPOINT
+) -> SAM2ImagePredictor:
+    model = build_sam2(config, checkpoint, device=device)
+    return SAM2ImagePredictor(sam_model=model)
+
+
+def load_sam_video_model(
+    device: torch.device,
+    config: str = SAM_CONFIG,
+    checkpoint: str = SAM_CHECKPOINT
+) -> Any:
+    return build_sam2_video_predictor(config, checkpoint, device=device)
+
+
+def run_sam_inference(
+    model: Any,
+    image: Image,
+    detections: sv.Detections
+) -> sv.Detections:
+    image = np.array(image.convert("RGB"))
+    model.set_image(image)
+    mask, score, _ = model.predict(box=detections.xyxy, multimask_output=False)
+
+    # dirty fix; remove this later
+    if len(mask.shape) == 4:
+        mask = np.squeeze(mask)
+
+    detections.mask = mask.astype(bool)
+    return detections