revert: to the b814393 and upgrade guidance_scale=3.8 and strength=0.6

Dockerfile

@@ -8,6 +8,7 @@ FROM nvidia/cuda:12.1.1-devel-ubuntu22.04
 # ============================================
 ENV PIP_CACHE_DIR=/data/.cache/pip \
     HF_HOME=/data/.cache/huggingface \
+    OMP_NUM_THREADS=4 \
     MAX_JOBS=1 \
     FORCE_CUDA=1 \
     PYTHONUNBUFFERED=1 \
@@ -40,11 +41,6 @@ WORKDIR /app
 RUN python3 -m pip install --upgrade pip setuptools wheel
 RUN pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
 
-# ============================================
-# Установка xformers для оптимизации
-# ============================================
-RUN pip install --no-cache-dir xformers==0.0.22.post7
-
 # ============================================
 # Установка зависимостей Python
 # ============================================

logic.py

@@ -5,148 +5,172 @@ from PIL import Image, ImageDraw
 import torch
 from loguru import logger
 import time
-
-# Импорты для iopaint (LaMa)
-from iopaint.model_manager import ModelManager
-from iopaint.schema import InpaintRequest, HDStrategy, LDMSampler
-
-# Импорты для diffusers (ControlNet)
-from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
+from diffusers import AutoPipelineForInpainting
 
 class WatermarkRemover:
     def __init__(self, device=None):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         logger.info(f"Using device: {self.device}")
 
+        # Lazy-loaded models
         self.detector = None
-        self.lama_model = None # <--- Модель из iopaint
-        self.controlnet_pipe = None
-        self.controlnet = None
-
-    def _prepare_image_for_diffusion(self, image: Image.Image, mask: Image.Image) -> Image.Image:
-        image_np = np.array(image.convert("RGB"))
-        mask_np = np.array(mask.convert("L"))
-        inpainted_np = cv2.inpaint(image_np, mask_np, inpaintRadius=15, flags=cv2.INPAINT_NS)
-        return Image.fromarray(inpainted_np)
+        self.inpainting_pipe = None
 
+    # ======================================================
+    # Lazy-load YOLO
+    # ======================================================
     def _load_detector(self):
         if self.detector is None:
             logger.info("Loading YOLOv8 custom model ('best.pt')...")
-            self.detector = YOLO("best.pt").to(self.device)
-            try: self.detector.fuse()
-            except Exception: pass
+            self.detector = YOLO("best.pt")
+            self.detector.to(self.device)
+            try:
+                self.detector.fuse()
+            except Exception:
+                pass
             logger.success("YOLOv8 model loaded successfully.")
 
-    def _load_lama_model(self):
-        # Инициализируем LaMa через iopaint
-        if self.lama_model is None:
-            logger.info("Initializing LaMa model via iopaint...")
-            self.lama_model = ModelManager(name="lama", device=self.device)
-            logger.success("LaMa model initialized successfully.")
-    
-    def _load_controlnet_model(self):
-        if self.controlnet_pipe is None:
-            logger.info("Loading ControlNet-Inpaint model...")
-            self.controlnet = ControlNetModel.from_pretrained(
-                "lllyasviel/control_v11p_sd15_inpaint",
-                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
-            )
-            self.controlnet_pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-                "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet,
+    # ======================================================
+    # Lazy-load Stable Diffusion
+    # ======================================================
+    def _load_inpainting_model(self):
+        if self.inpainting_pipe is None:
+            logger.info("Loading Stable Diffusion 1.5 Inpainting from Hugging Face Hub...")
+            # Возвращаемся к скачиванию из интернета
+            self.inpainting_pipe = AutoPipelineForInpainting.from_pretrained(
+                "runwayml/stable-diffusion-inpainting",
                 torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
-                safety_checker=None
+                safety_checker=None,
             ).to(self.device)
-            self.controlnet_pipe.scheduler = UniPCMultepScheduler.from_config(self.controlnet_pipe.scheduler.config)
-            try: self.controlnet_pipe.enable_xformers_memory_efficient_attention()
-            except Exception: pass
-            logger.success("ControlNet-Inpaint model loaded successfully.")
+            try:
+                self.inpainting_pipe.enable_attention_slicing()
+            except Exception:
+                pass
+            logger.success("Stable Diffusion 1.5 Inpainting model loaded successfully.")
 
+    # ======================================================
+    # Mask generation via YOLO
+    # ======================================================
     def _get_mask_yolo(self, image: Image.Image) -> Image.Image:
-        self._load_detector()
+        self._load_detector()  # ensure YOLO loaded
+
         img_np = np.array(image.convert("RGB"))
         results = self.detector.predict(img_np, conf=0.25, imgsz=864, device=self.device)
         mask = Image.new("L", image.size, 0)
+
         if results and len(results[0].boxes) > 0:
             draw = ImageDraw.Draw(mask)
             boxes = results[0].boxes.xyxy.cpu().numpy()
-            for bbox in boxes: draw.rectangle(list(bbox), fill=255)
+            logger.info(f"YOLO found {len(boxes)} watermark box(es).")
+            for bbox in boxes:
+                draw.rectangle(list(bbox), fill=255)
+        else:
+            logger.warning("No watermark detected.")
         return mask
 
+    # ======================================================
+    # Partial inpainting
+    # ======================================================
     def _inpaint_image(self, image: Image.Image, mask: Image.Image) -> Image.Image:
-        # ... код до ЭТАПА 1 без изменений
-        prompt = "ultra realistic photo of architecture, high detail, 4k, professional photography"
-        negative_prompt = "blurry, distorted, deformed, low quality, noise, grain, text, logo"
+        self._load_inpainting_model()  # ensure pipeline loaded
+
+        prompt = (
+            "ultra realistic photo of interior or exterior architecture, "
+            "natural lighting, clean surface, consistent material texture, realistic color balance"
+        )
+        negative_prompt = (
+            "text, logo, watermark, signature, fake object, furniture, table, person, "
+            "painting, mirror artifact, blurry, distorted, deformed, low quality, noise, grain"
+        )
+
+        logger.info("Running partial Stable Diffusion inpainting...")
+
         orig_w, orig_h = image.size
         mask_np = np.array(mask)
         ys, xs = np.where(mask_np > 0)
-        if xs.size == 0 or ys.size == 0: return image
+
+        if xs.size == 0 or ys.size == 0:
+            logger.info("Mask empty — skipping inpainting.")
+            return image
+
         pad = max(48, int(min(orig_w, orig_h) * 0.03))
-        x_min, y_min = max(int(xs.min()) - pad, 0), max(int(ys.min()) - pad, 0)
-        x_max, y_max = min(int(xs.max()) + pad, orig_w), min(int(ys.max()) + pad, orig_h)
+        x_min = max(int(xs.min()) - pad, 0)
+        x_max = min(int(xs.max()) + pad, orig_w)
+        y_min = max(int(ys.min()) - pad, 0)
+        y_max = min(int(ys.max()) + pad, orig_h)
+
         crop_box = (x_min, y_min, x_max, y_max)
-        original_crop = image.crop(crop_box)
-        mask_crop = mask.crop(crop_box)
-
-        # --- ЭТАП 1: Быстрая очистка с помощью iopaint (LaMa) ---
-        logger.info("Step 1: Pre-cleaning with iopaint/LaMa to create a clean reference...")
-        self._load_lama_model()
-        
-        # Конвертируем PIL в NumPy
-        original_crop_np = np.array(original_crop)
-        mask_crop_np = np.array(mask_crop.convert("L"))
-
-        # Создаем конфиг для iopaint
-        config = InpaintRequest(hd_strategy=HDStrategy.ORIGINAL)
-        
-        # Вызываем модель. Важно: iopaint возвращает BGR!
-        cleaned_bgr_np = self.lama_model(original_crop_np, mask_crop_np, config)
-        
-        # Конвертируем BGR -> RGB
-        cleaned_rgb_np = cv2.cvtColor(cleaned_bgr_np, cv2.COLOR_BGR2RGB)
-        
-        cleaned_reference_crop = Image.fromarray(cleaned_rgb_np)
-        logger.success("Clean reference created.")
-        
-        if self.device == "cuda": torch.cuda.empty_cache()
-
-        # --- ЭТАП 2: Качественная реставрация с помощью ControlNet (без изменений) ---
-        logger.info("Step 2: High-quality restoration with ControlNet-Inpaint...")
-        self._load_controlnet_model()
-        canvas_crop = self._prepare_image_for_diffusion(original_crop, mask_crop)
-        crop_w, crop_h = original_crop.size
-        new_w, new_h = int(np.ceil(crop_w / 8) * 8), int(np.ceil(crop_h / 8) * 8)
-        resized_canvas = canvas_crop.resize((new_w, new_h), resample=Image.LANCZOS)
-        resized_mask = mask_crop.resize((new_w, new_h), resample=Image.NEAREST)
-        resized_cleaned_reference = cleaned_reference_crop.resize((new_w, new_h), resample=Image.LANCZOS)
+        crop_img = image.crop(crop_box)
+        crop_mask = mask.crop(crop_box)
+
+        crop_w, crop_h = crop_img.size
+        max_side = 1024
+        scale = 1.0
+        if max(crop_w, crop_h) > max_side:
+            scale = max_side / max(crop_w, crop_h)
+
+        new_w = int(np.ceil((crop_w * scale) / 8) * 8)
+        new_h = int(np.ceil((crop_h * scale) / 8) * 8)
+
+        if (new_w, new_h) != (crop_w, crop_h):
+            resized_img = crop_img.resize((new_w, new_h), resample=Image.LANCZOS)
+            resized_mask = crop_mask.resize((new_w, new_h), resample=Image.LANCZOS)
+        else:
+            resized_img, resized_mask = crop_img, crop_mask
+
+        resized_mask = resized_mask.convert("L")
+        mask_thr = np.array(resized_mask)
+        mask_thr = (mask_thr > 127).astype(np.uint8) * 255
+        resized_mask = Image.fromarray(mask_thr, mode="L")
+
         with torch.inference_mode():
-            result = self.controlnet_pipe(
-                prompt=prompt, negative_prompt=negative_prompt, image=resized_canvas,
-                mask_image=resized_mask, control_image=resized_cleaned_reference,
-                num_inference_steps=25, guidance_scale=7.5,
-                controlnet_conditioning_scale=1.0, strength=1.0
+            result = self.inpainting_pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                image=resized_img,
+                mask_image=resized_mask,
+                num_inference_steps=27,
+                guidance_scale=3.8,
+                strength=0.6,
             ).images[0]
-        logger.success("High-quality restoration complete.")
-        
-        if result.size != original_crop.size: result_resized = result.resize(original_crop.size, resample=Image.LANCZOS)
-        else: result_resized = result
+
+        if result.size != crop_img.size:
+            result_resized = result.resize(crop_img.size, resample=Image.LANCZOS)
+        else:
+            result_resized = result
+
         base = image.copy()
-        base.paste(result_resized, (x_min, y_min), mask=mask_crop)
-        if self.device == "cuda": torch.cuda.empty_cache()
+        paste_mask = crop_mask.convert("L")
+        paste_mask = Image.fromarray((np.array(paste_mask) > 127).astype(np.uint8) * 255, mode="L")
+        base.paste(result_resized, (x_min, y_min), mask=paste_mask)
+
+        if self.device == "cuda":
+            torch.cuda.empty_cache()
+
         return base
 
+    # ======================================================
+    # Main process
+    # ======================================================
     def run(self, image: Image.Image) -> Image.Image:
         start_time = time.time()
         logger.info("Starting watermark removal...")
+
         mask_image = self._get_mask_yolo(image)
         mask_np = np.array(mask_image)
-        if not np.any(mask_np): return image
-        logger.info("Post-processing mask...")
+
+        if not np.any(mask_np):
+            logger.info("No watermark found. Returning original image.")
+            return image
+
+        logger.info("Post-processing mask (morphology)...")
         kernel = np.ones((15, 15), np.uint8)
         closed_mask = cv2.morphologyEx(mask_np, cv2.MORPH_CLOSE, kernel)
         final_kernel = np.ones((7, 7), np.uint8)
         processed_mask_np = cv2.dilate(closed_mask, final_kernel, iterations=1)
         processed_mask_pil = Image.fromarray(processed_mask_np)
         logger.success("Mask processed.")
+
         result_img = self._inpaint_image(image, processed_mask_pil)
         end_time = time.time()
         logger.success(f"Watermark removal completed in {end_time - start_time:.2f}s.")

requirements.txt

@@ -14,20 +14,22 @@ ultralytics==8.2.2
 # ===============================
 # Утилиты
 # ===============================
-numpy==1.26.4
+numpy
 loguru==0.7.2
 opencv-python-headless==4.9.0.80
 Pillow==9.5.0
 einops
 
 # ===============================
-# --- Инпеинтинг ---
+# --- Hugging Face Ecosystem ---
 # ===============================
-# Устанавливаем iopaint с "extras", чтобы он подтянул все AI-модели (LaMa, SD и т.д.)
-# Это также автоматически установит ПРАВИЛЬНЫЕ версии diffusers, transformers и т.д.
-iopaint[models]==1.2.1
+# Совместимые версии для SD2 Inpainting и torch 2.3+/CUDA 12.1
+diffusers==0.21.4
+transformers==4.30.2
+accelerate==0.21.0
+huggingface-hub==0.21.4
 
 # ===============================
-# Оптимизация
+# Прочее (иногда нужно diffusers)
 # ===============================
-xformers==0.0.22.post7
+safetensors