init commit

app.py

@@ -22,6 +22,7 @@ import json
 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 supervision as sv
 
 
 
@@ -41,6 +42,10 @@ taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).
 good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
 pipe = FluxInpaintPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
 
+FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=device)
+SAM_IMAGE_MODEL = load_sam_image_model(device=device)
+
+
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -129,7 +134,6 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
     return image_file
 
 
-@spaces.GPU(duration=50)
 def run_flux(
     image: Image.Image,
     mask: Image.Image,
@@ -154,28 +158,74 @@ def run_flux(
         seed_slicer = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed_slicer)
 
-    return PIPE(
-        prompt=prompt,
-        image=image,
-        mask_image=mask,
-        width=width,
-        height=height,
-        strength=strength_slider,
-        generator=generator,
-        num_inference_steps=num_inference_steps_slider,
-        max_sequence_length=256,
-        joint_attention_kwargs={"scale": lora_scale},
-    ).images[0]
-
-
-@spaces.GPU(duration=50)
+    with calculateDuration("run pipe"):
+        genearte_image = PIPE(
+            prompt=prompt,
+            image=image,
+            mask_image=mask,
+            width=width,
+            height=height,
+            strength=strength_slider,
+            generator=generator,
+            num_inference_steps=num_inference_steps_slider,
+            max_sequence_length=256,
+            joint_attention_kwargs={"scale": lora_scale},
+        ).images[0]
+    
+    return genearte_image
+
+
 def genearte_mask(image: Image.Image, masking_prompt_text: str) -> Image.Image:
     # generate mask by florence & sam
     print("Generating mask...")
-
-    return 
+    task_prompt = "<CAPTION_TO_PHRASE_GROUNDING>"
+    
+    with calculateDuration("FLORENCE"):
+        print(task_prompt, masking_prompt_text)
+        _, result = run_florence_inference(
+            model=FLORENCE_MODEL,
+            processor=FLORENCE_PROCESSOR,
+            device=device,
+            image=image,
+            task=task_prompt,
+            text=masking_prompt_text
+        )
+
+    with calculateDuration("sv.Detections"):
+        # start to dectect
+        detections = sv.Detections.from_lmm(
+            lmm=sv.LMM.FLORENCE_2,
+            result=result,
+            resolution_wh=image_input.size
+        )
+    
+    images = []
+
+    with calculateDuration("generate segmenet mask"):
+        # using sam generate segments images        
+        detections = run_sam_inference(SAM_IMAGE_MODEL, image, detections)
+        if len(detections) == 0:
+            gr.Info("No objects detected.")
+            return None
+        print("mask generated:", len(detections.mask))
+        kernel_size = dilate
+        kernel = np.ones((kernel_size, kernel_size), np.uint8)
+
+        for i in range(len(detections.mask)):
+            mask = detections.mask[i].astype(np.uint8) * 255
+            images.append(mask)
+
+        # merge mark into on image
+        merged_mask = np.zeros_like(images[0], dtype=np.uint8)
+        for mask in images:
+            merged_mask = cv2.bitwise_or(merged_mask, mask)
+        
+        images = [merged_mask]
+    
+    return images[0]
 
 
+@spaces.GPU(duration=120)
 def process(
     image_url: str,
     inpainting_prompt_text: str,
@@ -199,31 +249,32 @@ def process(
     if not image_url:
         gr.Info("please enter image url for inpaiting")
         result["message"] = "invalid image url"
-        return None, None, json.dumps(result)
+        return json.dumps(result)
 
     if not inpainting_prompt_text:
         gr.Info("Please enter inpainting text prompt.")
         result["message"] = "invalid inpainting prompt"
-        return None, None, json.dumps(result)
+        return json.dumps(result)
     
     if not masking_prompt_text:
         gr.Info("Please enter masking_prompt_text.")
         result["message"] = "invalid masking prompt"
-        return None, None, json.dumps(result)
-
+        return json.dumps(result)
 
-    image = load_image(image_url)
+    with calculateDuration("load image"):
+        image = load_image(image_url)
+    
     mask = genearte_mask(image, masking_prompt_text)
     
     if not image:
         gr.Info("Please upload an image.")
         result["message"] = "can not load image"
-        return None, None, json.dumps(result)
+        return json.dumps(result)
 
     if is_mask_empty(mask):
         gr.Info("Please draw a mask or enter a masking prompt.")
         result["message"] = "can not generate mask"
-        return None, None,  json.dumps(result)
+        return json.dumps(result)
 
     # generate
     width, height = calculate_image_dimensions_for_flux(original_resolution_wh=image.size)
@@ -243,14 +294,14 @@ def process(
         num_inference_steps_slider=num_inference_steps_slider,
         resolution_wh=(width, height)
     )
+
     if upload_to_r2:
         url = upload_image_to_r2(image, account_id, access_key, secret_key, bucket)
         result = {"status": "success", "url": url}
     else:
         result = {"status": "success", "message": "Image generated but not uploaded"}
 
-    return image, mask, json.dumps(result)
-
+    return json.dumps(result)
 
 
 with gr.Blocks() as demo:
@@ -309,11 +360,8 @@ with gr.Blocks() as demo:
                     value=0.9,
                 )
                 
-                
-
             with gr.Accordion("Advanced Settings", open=False):
                
-
                 with gr.Row():
                     mask_inflation_slider_component = gr.Slider(
                         label="Mask inflation",
@@ -370,23 +418,17 @@ with gr.Blocks() as demo:
                     )
                 
                 upload_to_r2 = gr.Checkbox(label="Upload to R2", value=False)
-                account_id = gr.Textbox(label="Account Id", placeholder="Enter R2 account id")
-                access_key = gr.Textbox(label="Access Key", placeholder="Enter R2 access key here")
-                secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here")
-                bucket = gr.Textbox(label="Bucket Name", placeholder="Enter R2 bucket name here")
+                with gr.Row():
+                    account_id = gr.Textbox(label="Account Id", placeholder="Enter R2 account id")
+                    bucket = gr.Textbox(label="Bucket Name", placeholder="Enter R2 bucket name here")
+
+                with gr.Row():
+                    access_key = gr.Textbox(label="Access Key", placeholder="Enter R2 access key here")
+                    secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here")
         
                 
         with gr.Column():
 
-            output_image_component = gr.Image(
-                type='pil', image_mode='RGB', label='Generated image', format="png")
-            
-            
-
-            with gr.Accordion("Debug", open=False):
-                output_mask_component = gr.Image(
-                    type='pil', image_mode='RGB', label='Input mask', format="png")
-
             output_json_component = gr.Textbox()
 
     submit_button_component.click(
@@ -411,8 +453,6 @@ with gr.Blocks() as demo:
             bucket
         ],
         outputs=[
-            output_image_component,
-            output_mask_component,
             output_json_component
         ]
     )