Spaces:
Configuration error
Configuration error
| """CtrlRegen engine — orchestrates the full watermark removal pipeline. | |
| Loads the base SD 1.5 model with a ControlNet (spatial control from | |
| canny edges) and a DINOv2-based IP Adapter (semantic control), then | |
| runs controllable regeneration with optional color matching. | |
| Attribution: | |
| Based on https://github.com/yepengliu/CtrlRegen (Apache-2.0). | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| import sys | |
| import time | |
| from typing import Any, Callable | |
| import torch | |
| from PIL import Image | |
| from progress import make_pipeline_progress | |
| logger = logging.getLogger(__name__) | |
| # --------------------------------------------------------------------------- | |
| # Availability checks — these imports are optional. | |
| # --------------------------------------------------------------------------- | |
| _HAS_CONTROLNET_AUX = False | |
| _HAS_COLOR_MATCHER = False | |
| _HAS_DIFFUSERS = False | |
| try: | |
| from diffusers import AutoencoderKL, ControlNetModel, UniPCMultistepScheduler | |
| from ctrlregen.pipeline import CustomCtrlRegenPipeline | |
| _HAS_DIFFUSERS = True | |
| except ImportError: | |
| AutoencoderKL = None # type: ignore[assignment,misc] | |
| ControlNetModel = None # type: ignore[assignment,misc] | |
| UniPCMultistepScheduler = None # type: ignore[assignment,misc] | |
| CustomCtrlRegenPipeline = None # type: ignore[assignment,misc] | |
| try: | |
| from controlnet_aux import CannyDetector | |
| _HAS_CONTROLNET_AUX = True | |
| except ImportError: | |
| CannyDetector = None # type: ignore[assignment,misc] | |
| try: | |
| from ctrlregen.color import color_match | |
| _HAS_COLOR_MATCHER = True | |
| except ImportError: | |
| color_match = None # type: ignore[assignment] | |
| CTRLREGEN_HF_REPO = "yepengliu/ctrlregen" | |
| SPATIAL_SUBFOLDER = "spatialnet_ckp/spatial_control_ckp_14000" | |
| SEMANTIC_SUBFOLDER = "semanticnet_ckp/models" | |
| SEMANTIC_WEIGHT_NAME = "semantic_control_ckp_435000.bin" | |
| DEFAULT_BASE_MODEL = "SG161222/Realistic_Vision_V4.0_noVAE" | |
| CUSTOM_VAE_ID = "stabilityai/sd-vae-ft-mse" | |
| PROCESS_SIZE = 512 | |
| DEFAULT_GUIDANCE_SCALE = 2.0 | |
| QUALITY_PROMPT = "best quality, high quality" | |
| NEGATIVE_PROMPT = "monochrome, lowres, bad anatomy, worst quality, low quality" | |
| CANNY_LOW_THRESHOLD = 100 | |
| CANNY_HIGH_THRESHOLD = 150 | |
| TILE_SIZE = 512 | |
| TILE_OVERLAP = 192 | |
| def is_ctrlregen_available() -> bool: | |
| """Return True when all CtrlRegen-specific dependencies are installed.""" | |
| return _HAS_DIFFUSERS and _HAS_CONTROLNET_AUX and _HAS_COLOR_MATCHER | |
| class CtrlRegenEngine: | |
| """End-to-end CtrlRegen watermark removal engine. | |
| Handles model loading, canny edge extraction, controlled denoising, | |
| and color-matched post-processing in a single ``run()`` call. | |
| """ | |
| def __init__( | |
| self, | |
| base_model_id: str | None = None, | |
| device: str = "cpu", | |
| torch_dtype: torch.dtype | None = None, | |
| hf_token: str | None = None, | |
| progress_callback: Callable[[str], None] | None = None, | |
| ) -> None: | |
| if not is_ctrlregen_available(): | |
| missing: list[str] = [] | |
| if not _HAS_DIFFUSERS: | |
| missing.extend(["diffusers", "transformers", "accelerate"]) | |
| if not _HAS_CONTROLNET_AUX: | |
| missing.append("controlnet-aux") | |
| if not _HAS_COLOR_MATCHER: | |
| missing.append("color-matcher") | |
| logger.info("Auto-installing missing dependencies: %s", missing) | |
| import subprocess | |
| try: | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", *missing], | |
| stdout=subprocess.DEVNULL, | |
| stderr=subprocess.DEVNULL, | |
| ) | |
| except (subprocess.CalledProcessError, FileNotFoundError): | |
| raise ImportError( | |
| "Failed to auto-install missing dependencies: " | |
| + ", ".join(missing) | |
| + ". Try manually: pip install --force-reinstall noai-watermark" | |
| ) | |
| self.base_model_id = base_model_id or DEFAULT_BASE_MODEL | |
| self.device = device | |
| self.torch_dtype = torch_dtype or ( | |
| torch.float16 if device != "cpu" else torch.float32 | |
| ) | |
| self.hf_token: str | None = hf_token or os.environ.get("HF_TOKEN") | |
| self._progress_callback = progress_callback | |
| self._pipeline: CustomCtrlRegenPipeline | None = None # type: ignore[assignment] | |
| self._canny_detector: CannyDetector | None = None # type: ignore[assignment] | |
| def _set_progress(self, message: str) -> None: | |
| if self._progress_callback is None: | |
| return | |
| try: | |
| self._progress_callback(message) | |
| except Exception: | |
| pass | |
| # ------------------------------------------------------------------ | |
| # Loading | |
| # ------------------------------------------------------------------ | |
| def load(self) -> None: | |
| """Download and assemble the full CtrlRegen pipeline.""" | |
| if self._pipeline is not None: | |
| return | |
| token_kwargs: dict[str, Any] = {} | |
| if self.hf_token: | |
| token_kwargs["token"] = self.hf_token | |
| self._set_progress( | |
| f"Loading CtrlRegen spatial ControlNet from {CTRLREGEN_HF_REPO}..." | |
| ) | |
| logger.info("Loading ControlNet from %s/%s", CTRLREGEN_HF_REPO, SPATIAL_SUBFOLDER) | |
| controlnet = [ControlNetModel.from_pretrained( | |
| CTRLREGEN_HF_REPO, | |
| subfolder=SPATIAL_SUBFOLDER, | |
| torch_dtype=self.torch_dtype, | |
| **token_kwargs, | |
| )] | |
| self._set_progress( | |
| f"Loading SD base model ({self.base_model_id}) for CtrlRegen pipeline..." | |
| ) | |
| logger.info("Loading base pipeline from %s", self.base_model_id) | |
| pipe = CustomCtrlRegenPipeline.from_pretrained( | |
| self.base_model_id, | |
| controlnet=controlnet, | |
| torch_dtype=self.torch_dtype, | |
| safety_checker=None, | |
| requires_safety_checker=False, | |
| **token_kwargs, | |
| ) | |
| self._set_progress( | |
| f"Loading CtrlRegen semantic IP-Adapter + DINOv2 from {CTRLREGEN_HF_REPO}..." | |
| ) | |
| logger.info("Loading IP-Adapter from %s/%s", CTRLREGEN_HF_REPO, SEMANTIC_SUBFOLDER) | |
| pipe.load_ctrlregen_ip_adapter( | |
| CTRLREGEN_HF_REPO, | |
| subfolder=SEMANTIC_SUBFOLDER, | |
| weight_name=SEMANTIC_WEIGHT_NAME, | |
| **token_kwargs, | |
| ) | |
| from transformers import AutoImageProcessor, AutoModel | |
| pipe.image_encoder = AutoModel.from_pretrained("facebook/dinov2-giant").to( | |
| self.device, dtype=self.torch_dtype | |
| ) | |
| pipe.feature_extractor = AutoImageProcessor.from_pretrained("facebook/dinov2-giant") | |
| self._set_progress(f"Loading custom VAE ({CUSTOM_VAE_ID})...") | |
| logger.info("Loading VAE from %s", CUSTOM_VAE_ID) | |
| pipe.vae = AutoencoderKL.from_pretrained( | |
| CUSTOM_VAE_ID, torch_dtype=self.torch_dtype, **token_kwargs, | |
| ) | |
| self._set_progress("Configuring UniPC scheduler...") | |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
| pipe.set_ip_adapter_scale(1.0) | |
| self._set_progress(f"Moving CtrlRegen pipeline to {self.device}...") | |
| pipe = pipe.to(self.device) | |
| if hasattr(pipe, "enable_xformers_memory_efficient_attention"): | |
| try: | |
| pipe.enable_xformers_memory_efficient_attention() | |
| except Exception: | |
| pass | |
| self._pipeline = pipe | |
| self._canny_detector = CannyDetector() | |
| self._set_progress("CtrlRegen pipeline ready.") | |
| logger.info("CtrlRegen pipeline loaded on %s", self.device) | |
| # ------------------------------------------------------------------ | |
| # Inference — public entry point | |
| # ------------------------------------------------------------------ | |
| def run( | |
| self, | |
| image: Image.Image, | |
| strength: float = 0.5, | |
| num_inference_steps: int = 50, | |
| guidance_scale: float = DEFAULT_GUIDANCE_SCALE, | |
| seed: int | None = None, | |
| ) -> Image.Image: | |
| """Run CtrlRegen watermark removal on a single image. | |
| Images that fit within ``TILE_SIZE`` (512) are processed as a | |
| single pass. Larger images are split into overlapping tiles. | |
| """ | |
| self.load() | |
| assert self._pipeline is not None | |
| assert self._canny_detector is not None | |
| orig_w, orig_h = image.size | |
| orig_image = image | |
| t0 = time.monotonic() | |
| needs_tiling = orig_w > TILE_SIZE or orig_h > TILE_SIZE | |
| if needs_tiling: | |
| from ctrlregen.tiling import resize_center_crop, run_tiled | |
| aligned_w = orig_w // 8 * 8 | |
| aligned_h = orig_h // 8 * 8 | |
| if aligned_w != orig_w or aligned_h != orig_h: | |
| image = image.resize((aligned_w, aligned_h), Image.LANCZOS) | |
| regen_image = run_tiled( | |
| pipeline=self._pipeline, | |
| canny_detector=self._canny_detector, | |
| image=image, | |
| strength=strength, | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale, | |
| seed=seed, | |
| tile_size=TILE_SIZE, | |
| tile_overlap=TILE_OVERLAP, | |
| quality_prompt=QUALITY_PROMPT, | |
| negative_prompt=NEGATIVE_PROMPT, | |
| canny_low=CANNY_LOW_THRESHOLD, | |
| canny_high=CANNY_HIGH_THRESHOLD, | |
| device=self.device, | |
| set_progress=self._set_progress, | |
| ip_adapter_image=orig_image, | |
| ) | |
| else: | |
| from ctrlregen.tiling import resize_center_crop | |
| proc_image = resize_center_crop(image, PROCESS_SIZE) | |
| self._set_progress( | |
| f"Preprocessed {orig_w}x{orig_h}px → " | |
| f"{proc_image.size[0]}x{proc_image.size[1]}px" | |
| ) | |
| regen_image = self._run_single( | |
| proc_image, strength, num_inference_steps, guidance_scale, seed, | |
| ) | |
| if regen_image.size != (orig_w, orig_h): | |
| self._set_progress( | |
| f"Resizing {regen_image.size[0]}x{regen_image.size[1]}px " | |
| f"→ {orig_w}x{orig_h}px..." | |
| ) | |
| regen_image = regen_image.resize((orig_w, orig_h), Image.LANCZOS) | |
| self._set_progress(f"Applying color matching at {orig_w}x{orig_h}px...") | |
| output = color_match(reference=orig_image, source=regen_image) | |
| self._set_progress( | |
| f"✓ CtrlRegen done · {orig_w}x{orig_h}px · " | |
| f"{time.monotonic() - t0:.0f}s total" | |
| ) | |
| return output | |
| # ------------------------------------------------------------------ | |
| # Single-image path (image <= 512x512) | |
| # ------------------------------------------------------------------ | |
| def _run_single( | |
| self, | |
| image: Image.Image, | |
| strength: float, | |
| num_inference_steps: int, | |
| guidance_scale: float, | |
| seed: int | None, | |
| ) -> Image.Image: | |
| """Process a single 512x512 image through the CtrlRegen pipeline.""" | |
| w, h = image.size | |
| effective_steps = max(1, int(num_inference_steps * strength)) | |
| self._set_progress(f"Extracting canny edges ({w}x{h}px, " | |
| f"thresholds {CANNY_LOW_THRESHOLD}/{CANNY_HIGH_THRESHOLD})...") | |
| control_image = self._canny_detector( | |
| image, low_threshold=CANNY_LOW_THRESHOLD, high_threshold=CANNY_HIGH_THRESHOLD, | |
| ) | |
| generator = torch.manual_seed(seed if seed is not None else 0) | |
| self._set_progress( | |
| f"Config: strength={strength}, steps={num_inference_steps} " | |
| f"(~{effective_steps} effective), guidance={guidance_scale}" | |
| ) | |
| step_cb, first_step, pipeline_done, start_updater = make_pipeline_progress( | |
| effective_steps, self.device, self._set_progress, | |
| label="CtrlRegen denoising", | |
| ) | |
| start_updater() | |
| try: | |
| result = self._pipeline( | |
| prompt=QUALITY_PROMPT, | |
| negative_prompt=NEGATIVE_PROMPT, | |
| image=[image], | |
| control_image=[control_image], | |
| controlnet_conditioning_scale=1.0, | |
| ip_adapter_image=[image], | |
| strength=strength, | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator, | |
| control_guidance_start=0.0, | |
| control_guidance_end=1.0, | |
| callback=step_cb, | |
| callback_steps=1, | |
| ) | |
| except TypeError: | |
| first_step.set() | |
| result = self._pipeline( | |
| prompt=QUALITY_PROMPT, | |
| negative_prompt=NEGATIVE_PROMPT, | |
| image=[image], | |
| control_image=[control_image], | |
| controlnet_conditioning_scale=1.0, | |
| ip_adapter_image=[image], | |
| strength=strength, | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator, | |
| ) | |
| finally: | |
| first_step.set() | |
| pipeline_done.set() | |
| return result.images[0] | |