Sync from GitHub via hub-sync

Dockerfile

@@ -6,6 +6,7 @@ RUN apt-get update && apt-get install -y \
     git \
     libgl1 \
     libglib2.0-0 \
+    libgomp1 \
     libsm6 \
     libxext6 \
     libxrender-dev \
@@ -27,9 +28,13 @@ ENV HF_HUB_ENABLE_HF_TRANSFER=1 \
     CELERY_RESULT_BACKEND=redis://localhost:6379/1
 
 # Step 3. Download models for caching
+RUN mkdir -p ./models
 RUN python -c "from huggingface_hub import snapshot_download; snapshot_download('zeyadcode/ov_owlv2_model', local_dir='./models/ov_owlv2_model')"    
 RUN python -c "from ultralytics import SAM; import os; os.chdir('./models'); SAM('mobile_sam.pt')"
 RUN python -c "from transformers import AutoProcessor; AutoProcessor.from_pretrained('google/owlv2-base-patch16-ensemble')"
+RUN python -c "from simple_lama_inpainting import SimpleLama; SimpleLama()"
+RUN python -c "from optimum.intel import OVStableDiffusionXLInpaintPipeline; pipe = OVStableDiffusionXLInpaintPipeline.from_pretrained('stabilityai/sdxl-turbo', export=True, device='CPU'); pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1); pipe.save_pretrained('./models/ov_sdxl_turbo_inpaint')"
+
 
 # Step 4: Copy the rest of the application
 COPY . .

config.py

@@ -25,7 +25,10 @@ DETECTOR_OUT_PATH = "detector_boxes.json"
 
 INPAINTOR_OUT_PATH = "inpainted.png"
 
-INPAINTOR_MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"
+INPAINTOR_MODEL_NAME = "stabilityai/sdxl-turbo"
+INPAINTOR_MODEL_PATH = Path(os.getenv("INPAINTOR_MODEL_PATH", "models/ov_sdxl_turbo_inpaint"))
+INPAINTOR_DEVICE = os.getenv("INPAINTOR_DEVICE", "CPU")
+INPAINTOR_IMAGE_SIZE = 512
 
 CELERY_BROKER_URL = os.getenv("CELERY_BROKER_URL", "redis://localhost:6379/0")
 CELERY_RESULT_BACKEND = os.getenv("CELERY_RESULT_BACKEND", "redis://localhost:6379/1")

main.py

@@ -10,7 +10,7 @@ import shutil
 import tasks as tasks
 from celery.result import AsyncResult
 import config as config
-from schemas import SegmentRequest
+from schemas import InpaintRequest, SegmentRequest
 from services import job_lifecycle
 
 job_lifecycle.ensure_upload_dir()
@@ -85,9 +85,19 @@ async def segment_objects(job_id: str, data: SegmentRequest):
 
 
 @app.post("/process/inpaint/{job_id}")
-async def inapint(job_id, prompt: str = Form(...)):
+async def inpaint(job_id: str, data: InpaintRequest):
     tasks.celery_app.backend.delete(f"celery-task-meta-{job_id}")
-    tasks.inpaint.apply_async(args=[job_id, prompt], task_id=job_id)
+    tasks.inpaint.apply_async(
+        args=[
+            job_id,
+            data.mode,
+            data.positive_prompt,
+            data.negative_prompt,
+            data.num_inference_steps,
+        ],
+        task_id=job_id,
+    )
+    return job_id
 
 
 if __name__ == "__main__":

requirements.txt

@@ -10,13 +10,17 @@ opencv-python-headless
 Pillow
 
 # Hugging Face & Model Hub
+diffusers
 huggingface-hub
 hf-transfer
+safetensors
 transformers
 
 # Model Execution Frameworks
 openvino
 optimum
 optimum-intel
+torch
+torchvision
 ultralytics
-simple-lama-inpainting
+simple-lama-inpainting

schemas/__init__.py

@@ -1 +1,2 @@
-from .segment_request import SegmentRequest
+from .segment_request import SegmentRequest
+from .inpaint_request import InpaintRequest

schemas/inpaint_request.py

@@ -0,0 +1,10 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class InpaintRequest(BaseModel):
+    mode: Literal["blur", "remove", "diffusion"] = "blur"
+    positive_prompt: str = ""
+    negative_prompt: str = ""
+    num_inference_steps: int = Field(default=4, ge=4, le=15)

services/model_manager.py

@@ -3,6 +3,7 @@ from ultralytics import SAM
 import openvino as ov
 import config as config
 from optimum.intel import OVStableDiffusionXLInpaintPipeline
+from simple_lama_inpainting import SimpleLama
 
 
 def get_detector_model(MODELS: dict):
@@ -29,16 +30,30 @@ def get_segmentor_model(MODELS: dict):
 
 def get_inpaintor_model(MODELS: dict):
     if MODELS['inpaintor'] is None:
-        MODELS['inpaintor'] = OVStableDiffusionXLInpaintPipeline.from_pretrained(
-            config.INPAINTOR_MODEL_NAME,
-            device="CPU" # or "GPU"
+        model_source = (
+            config.INPAINTOR_MODEL_PATH
+            if config.INPAINTOR_MODEL_PATH.exists()
+            else config.INPAINTOR_MODEL_NAME
         )
+        kwargs = {"device": config.INPAINTOR_DEVICE}
+        if model_source == config.INPAINTOR_MODEL_NAME:
+            kwargs["export"] = True
+
+        pipe = OVStableDiffusionXLInpaintPipeline.from_pretrained(
+            str(model_source),
+            **kwargs,
+        )
+        pipe.reshape(
+            batch_size=1,
+            height=config.INPAINTOR_IMAGE_SIZE,
+            width=config.INPAINTOR_IMAGE_SIZE,
+            num_images_per_prompt=1,
+        )
+        pipe.compile()
+        MODELS['inpaintor'] = pipe
     return MODELS['inpaintor']
     
 def get_removing_model(MODELS: dict):
-    if MODELS['obj_remove'] is None:
-        MODELS['obj_remove'] = OVStableDiffusionXLInpaintPipeline.from_pretrained(
-            config.INPAINTOR_MODEL_NAME,
-            device="CPU" # or "GPU"
-        )
-    return MODELS['obj_remove']
+    if MODELS['remover'] is None:
+        MODELS['remover'] = SimpleLama()
+    return MODELS['remover']

tasks.py

@@ -15,6 +15,7 @@ MODELS: dict[str, Any] = {
     "detector_processor": None,
     "segmentor": None,
     "inpaintor": None,
+    "remover": None,
 }
 
 celery_app = Celery(
@@ -96,30 +97,46 @@ def segment_objects(self, job_id: str, bboxes=None, points=None, point_labels=No
 
 
 @celery_app.task(bind=True)
-def inpaint(self, job_id: str, prompt: str, negative_prompt: str = "", num_inference_steps: int = 6, guidance_scale: float = 0.0) -> bool:
+def inpaint(
+    self,
+    job_id: str,
+    mode: str = "blur",
+    positive_prompt: str = "",
+    negative_prompt: str = "",
+    num_inference_steps: int = 4,
+) -> bool:
     job_path = config.UPLOAD_DIR / job_id
+    save_path = job_path / config.INPAINTOR_OUT_PATH
+    if save_path.exists():
+        save_path.unlink()
+
     orig_img = Image.open(job_path / config.INPUT_IMG_NAME).convert("RGB")
     mask_img = Image.open(job_path / config.SEGMENTOR_OUT_BIN_PATH).convert("L")
 
-    mode = "blur"
-
-    if mode == "inpaint":
+    if mode == "diffusion":
         model = model_manager.get_inpaintor_model(MODELS)
-        inpainted_img = model(
-            prompt=prompt,
+        model_size = (config.INPAINTOR_IMAGE_SIZE, config.INPAINTOR_IMAGE_SIZE)
+        generated = model(
+            prompt=positive_prompt,
             negative_prompt=negative_prompt,
-            image=orig_img,
-            mask_image=mask_img,
-            num_inference_steps=6,
+            image=orig_img.resize(model_size),
+            mask_image=mask_img.resize(model_size, Image.Resampling.NEAREST),
+            num_inference_steps=num_inference_steps,
             guidance_scale=0.0,
         ).images[0]
+        result = generated.resize(orig_img.size)
 
     elif mode == "remove":
         model = model_manager.get_removing_model(MODELS)
+        result = model(orig_img, mask_img)
 
     elif mode == "blur":
         result = blur_img_with_mask(orig_img, mask_img)
-        result.save(job_path / config.INPAINTOR_OUT_PATH)
+
+    else:
+        raise ValueError(f"Unsupported inpaint mode: {mode}")
+
+    result.convert("RGB").save(save_path)
         
     if self.request.id is None:
         return result
@@ -131,18 +148,24 @@ def inpaint(self, job_id: str, prompt: str, negative_prompt: str = "", num_infer
 if __name__ == "__main__":
     import matplotlib.pyplot as plt
 
-    prompt = "hat"
-    img = Image.open(config.UPLOAD_DIR / config.DEBUG_JOB_ID / config.INPUT_IMG_NAME)
+    # prompt = "hat"
+    # img = Image.open(config.UPLOAD_DIR / config.DEBUG_JOB_ID / config.INPUT_IMG_NAME)
 
     # %%
-    detection_res = detect_objects(job_id=config.DEBUG_JOB_ID, prompt=prompt)
-    plot_groundingdino_boxes(img, detection_res)
+    # detection_res = detect_objects(job_id=config.DEBUG_JOB_ID, prompt=prompt)
+    # plot_groundingdino_boxes(img, detection_res)
 
     # %%
-    segment_res = segment_objects.run(job_id=config.DEBUG_JOB_ID, bboxes=detection_res["boxes"])
+    # segment_res = segment_objects.run(job_id=config.DEBUG_JOB_ID, bboxes=detection_res["boxes"])
 
     # %%
-    inpainted_res = inpaint(job_id=config.DEBUG_JOB_ID, prompt="red hat high quality", negative_prompt="blurry, messy, bad looking", num_inference_steps=4, guidance_scale=0)
+    inpainted_res = inpaint(
+        job_id=config.DEBUG_JOB_ID,
+        mode="diffusion",
+        positive_prompt="red hat high quality",
+        negative_prompt="blurry, messy, bad looking",
+        num_inference_steps=4,
+    )
 
     # %%
     img = segment_res.plot()