""" providers.py - Hybrid local editing providers for ImageEditter. This module exposes: - A direct classical CV provider. - A relaxed custom diffusion backend. - Optional foundation-model backends. - A hybrid provider that orchestrates everything through the editing stack. """ from __future__ import annotations import os from dataclasses import dataclass, field from typing import Any, Optional import numpy as np import torch from PIL import Image, ImageFilter from server.cv_engine import CVEngine, OperationContext, OperationStep from server.editing_stack import EditingOrchestrator, InstructionParserStage @dataclass class EditResult: image: Image.Image message: str provider: str used_fallback: bool = False steps: list[str] = field(default_factory=list) metadata: dict[str, Any] = field(default_factory=dict) class BaseEditProvider: provider_id = "base" provider_label = "Base provider" supports_broad_editing = False supports_sampling_controls = False supports_generation = False supports_inpaint = False supports_batch = False supports_style_transfer = False supports_background_ops = False supports_upscale = False supports_diffusion = False prompt_hint = "Describe your edit." model_loaded = True def warmup(self) -> bool: return self.model_loaded def edit( self, image: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, mask: Optional[Image.Image] = None, reference_image: Optional[Image.Image] = None, background_image: Optional[Image.Image] = None, ) -> EditResult: raise NotImplementedError def generate( self, prompt: str, width: int = 768, height: int = 768, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: raise NotImplementedError def inpaint( self, image: Image.Image, mask: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: raise NotImplementedError def batch_edit( self, images: list[Image.Image], prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> list[EditResult]: return [ self.edit( image=image, prompt=prompt, num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, ) for image in images ] def style_transfer( self, image: Image.Image, reference_image: Image.Image, prompt: str = "", seed: Optional[int] = None, ) -> EditResult: raise NotImplementedError def background_edit( self, image: Image.Image, prompt: str, background_image: Optional[Image.Image] = None, seed: Optional[int] = None, ) -> EditResult: raise NotImplementedError def upscale( self, image: Image.Image, scale: float = 2.0, prompt: str = "", seed: Optional[int] = None, ) -> EditResult: raise NotImplementedError def capabilities(self) -> list[dict[str, Any]]: return [] def presets(self) -> list[dict[str, Any]]: return [] class CVDirectProvider(BaseEditProvider): provider_id = "cv-engine" provider_label = "Local CV engine" supports_broad_editing = True supports_sampling_controls = False supports_generation = True supports_inpaint = True supports_batch = True supports_style_transfer = True supports_background_ops = True supports_upscale = True supports_diffusion = False prompt_hint = ( "Local classical editor with multi-step prompt routing, background replacement, " "style transfer, inpainting, relighting, portrait cleanup, and procedural generation." ) def __init__(self, engine: Optional[CVEngine] = None): self.engine = engine or CVEngine() self.parser = InstructionParserStage(self.engine) def capabilities(self) -> list[dict[str, Any]]: return self.engine.list_capabilities() def presets(self) -> list[dict[str, Any]]: return self.engine.list_presets() def edit( self, image: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, mask: Optional[Image.Image] = None, reference_image: Optional[Image.Image] = None, background_image: Optional[Image.Image] = None, ) -> EditResult: parsed = self.parser.parse(prompt) steps = parsed.cv_steps or [OperationStep(name="auto_enhance", params={"amount": 0.3})] context = OperationContext( mask=mask, reference_image=reference_image, background_image=background_image, prompt=prompt, seed=seed, ) result = self.engine.execute_pipeline(image=image.convert("RGB"), steps=steps, context=context) message = "CV engine edit complete." if parsed.expansion_notes: message += " Strategy: " + ", ".join(dict.fromkeys(parsed.expansion_notes)) + "." if result.traces: message += " Pipeline: " + ", ".join(trace.name for trace in result.traces) + "." return EditResult( image=result.image.convert("RGB"), message=message, provider=self.provider_id, used_fallback=False, steps=[trace.name for trace in result.traces], metadata=result.metadata, ) def generate( self, prompt: str, width: int = 768, height: int = 768, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: image = self.engine.procedural_generate(prompt=prompt, size=(width, height), seed=seed) parsed = self.parser.parse(prompt) steps = ["scene_generate"] if parsed.cv_steps: result = self.engine.execute_pipeline( image=image, steps=[step for step in parsed.cv_steps if step.name != "scene_generate"], context=OperationContext(prompt=prompt, seed=seed), ) image = result.image steps.extend(trace.name for trace in result.traces) return EditResult( image=image.convert("RGB"), message="CV engine generation complete. Pipeline: " + ", ".join(steps) + ".", provider=self.provider_id, used_fallback=False, steps=steps, metadata={"mode": "generate"}, ) def inpaint( self, image: Image.Image, mask: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: steps = [OperationStep(name="object_remove", params={"amount": 0.85})] if prompt.strip(): parsed = self.parser.parse(prompt) steps.extend(step for step in parsed.cv_steps if step.name != "object_remove") result = self.engine.execute_pipeline( image=image.convert("RGB"), steps=steps, context=OperationContext(mask=mask, prompt=prompt, seed=seed), ) return EditResult( image=result.image.convert("RGB"), message="CV inpainting complete. Pipeline: " + ", ".join(trace.name for trace in result.traces) + ".", provider=self.provider_id, used_fallback=False, steps=[trace.name for trace in result.traces], metadata={"mode": "inpaint"}, ) def style_transfer( self, image: Image.Image, reference_image: Image.Image, prompt: str = "", seed: Optional[int] = None, ) -> EditResult: steps = [OperationStep(name="style_reference", params={"amount": 0.8})] if prompt.strip(): parsed = self.parser.parse(prompt) steps.extend(step for step in parsed.cv_steps if step.name != "style_reference") result = self.engine.execute_pipeline( image=image.convert("RGB"), steps=steps, context=OperationContext(reference_image=reference_image, prompt=prompt, seed=seed), ) return EditResult( image=result.image.convert("RGB"), message="Local style transfer complete. Pipeline: " + ", ".join(trace.name for trace in result.traces) + ".", provider=self.provider_id, used_fallback=False, steps=[trace.name for trace in result.traces], metadata={"mode": "style-transfer"}, ) def background_edit( self, image: Image.Image, prompt: str, background_image: Optional[Image.Image] = None, seed: Optional[int] = None, ) -> EditResult: steps = [OperationStep(name="background_replace", params={"amount": 0.8})] parsed = self.parser.parse(prompt or "replace the background") steps.extend(step for step in parsed.cv_steps if step.name != "background_replace") result = self.engine.execute_pipeline( image=image.convert("RGB"), steps=steps, context=OperationContext(background_image=background_image, prompt=prompt, seed=seed), ) return EditResult( image=result.image.convert("RGB"), message="Background edit complete. Pipeline: " + ", ".join(trace.name for trace in result.traces) + ".", provider=self.provider_id, used_fallback=False, steps=[trace.name for trace in result.traces], metadata={"mode": "background"}, ) def upscale( self, image: Image.Image, scale: float = 2.0, prompt: str = "", seed: Optional[int] = None, ) -> EditResult: result = self.engine.apply_operation(image.convert("RGB"), "super_res", amount=0.8, scale=scale) result = self.engine.apply_operation(result, "clarity", amount=0.18) result = self.engine.apply_operation(result, "sharpen", amount=0.15) return EditResult( image=result.convert("RGB"), message=f"Upscale complete at {scale:.2f}x.", provider=self.provider_id, used_fallback=False, steps=["super_res", "clarity", "sharpen"], metadata={"mode": "upscale", "scale": scale}, ) class CustomDiffusionProvider(BaseEditProvider): provider_id = "custom-diffusion" provider_label = "Custom diffusion backend" supports_broad_editing = True supports_sampling_controls = True supports_generation = True supports_inpaint = True supports_diffusion = True prompt_hint = ( "Repo-native VAE/U-Net/CLIP diffusion backend. Best used inside the hybrid stack " "for semantic clothing, structure, and harmonization edits." ) def __init__(self, pipeline, fallback: Optional[CVDirectProvider] = None): self.pipeline = pipeline self.fallback = fallback def _get_steps(self, num_steps: int, default_min: int, default_max: int) -> int: if not torch.cuda.is_available(): # CPU mode: use fewer steps (e.g., 6 to 12 steps) return max(6, min(num_steps, 12)) return max(default_min, min(num_steps, default_max)) def edit( self, image: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, mask: Optional[Image.Image] = None, reference_image: Optional[Image.Image] = None, background_image: Optional[Image.Image] = None, ) -> EditResult: try: edited = self.pipeline.edit( image=image.convert("RGB"), prompt=prompt, num_steps=self._get_steps(num_steps, 12, 48), text_guidance_scale=float(np.clip(text_guidance_scale, 2.0, 10.0)), image_guidance_scale=float(np.clip(image_guidance_scale, 0.6, 3.0)), seed=seed, ).convert("RGB") except Exception as exc: if self.fallback is not None: fallback = self.fallback.edit(image=image, prompt=prompt, seed=seed) fallback.provider = self.provider_id fallback.used_fallback = True fallback.message = f"Custom diffusion failed and CV fallback was used: {exc}" return fallback raise if self._looks_degenerate(image, edited): if self.fallback is not None: fallback_res = self.fallback.edit( image=image, prompt=prompt, num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, mask=mask, reference_image=reference_image, background_image=background_image, ) fallback_res.provider = self.provider_id fallback_res.used_fallback = True fallback_res.message = "Local CV pipeline was automatically engaged to prevent custom generative noise patterns." return fallback_res stabilized = self._stabilize(image.convert("RGB"), edited) return EditResult( image=stabilized, message="Custom diffusion edit completed, but a stabilization blend was applied.", provider=self.provider_id, used_fallback=True, steps=["diffusion-stabilize"], ) return EditResult( image=edited, message="Custom diffusion edit complete.", provider=self.provider_id, used_fallback=False, steps=["diffusion"], ) def generate( self, prompt: str, width: int = 768, height: int = 768, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: try: image = self.pipeline.generate( prompt=prompt, num_steps=self._get_steps(num_steps, 16, 56), text_guidance_scale=float(np.clip(text_guidance_scale, 3.0, 11.0)), image_guidance_scale=max(0.0, image_guidance_scale * 0.1), seed=seed, width=width, height=height, ).convert("RGB") return EditResult( image=image, message="Custom diffusion generation complete.", provider=self.provider_id, used_fallback=False, steps=["diffusion-generate"], ) except Exception as exc: if self.fallback is not None: fallback = self.fallback.generate(prompt=prompt, width=width, height=height, seed=seed) fallback.provider = self.provider_id fallback.used_fallback = True fallback.message = f"Custom diffusion generation failed and CV generation was used: {exc}" return fallback raise def inpaint( self, image: Image.Image, mask: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: try: result = self.pipeline.inpaint( image=image.convert("RGB"), mask=mask.convert("L"), prompt=prompt, num_steps=self._get_steps(num_steps, 12, 42), text_guidance_scale=float(np.clip(text_guidance_scale, 3.0, 10.0)), image_guidance_scale=float(np.clip(image_guidance_scale, 0.8, 3.0)), seed=seed, ).convert("RGB") if self._looks_degenerate(image, result): if self.fallback is not None: fallback_res = self.fallback.inpaint(image=image, mask=mask, prompt=prompt, seed=seed) fallback_res.provider = self.provider_id fallback_res.used_fallback = True fallback_res.message = "Local CV inpainting was automatically engaged to prevent custom generative noise patterns." return fallback_res return EditResult( image=result, message="Custom diffusion inpainting complete.", provider=self.provider_id, used_fallback=False, steps=["diffusion-inpaint"], ) except Exception as exc: if self.fallback is not None: fallback = self.fallback.inpaint(image=image, mask=mask, prompt=prompt, seed=seed) fallback.provider = self.provider_id fallback.used_fallback = True fallback.message = f"Custom diffusion inpainting failed and CV inpainting was used: {exc}" return fallback raise def _looks_degenerate(self, input_image: Image.Image, output_image: Image.Image) -> bool: input_small = np.asarray(input_image.convert("RGB").resize((128, 128)), dtype=np.float32) / 255.0 output_small = np.asarray(output_image.convert("RGB").resize((128, 128)), dtype=np.float32) / 255.0 pixel_delta = float(np.mean(np.abs(output_small - input_small))) low_frequency_delta = float( np.mean( np.abs( np.asarray(input_image.convert("RGB").resize((128, 128)).filter(ImageFilter.GaussianBlur(radius=3)), dtype=np.float32) / 255.0 - np.asarray(output_image.convert("RGB").resize((128, 128)).filter(ImageFilter.GaussianBlur(radius=3)), dtype=np.float32) / 255.0 ) ) ) output_std = float(output_small.std()) return output_std < 0.02 or (pixel_delta > 0.62 and low_frequency_delta > 0.45) def _stabilize(self, input_image: Image.Image, output_image: Image.Image) -> Image.Image: stabilized = Image.blend(output_image.convert("RGB"), input_image.convert("RGB").resize(output_image.size), alpha=0.35) return stabilized.filter(ImageFilter.UnsharpMask(radius=1.2, percent=110, threshold=2)) class FoundationModelProvider(BaseEditProvider): backend_name = "foundation" model_id = "" requires_cuda = False supports_broad_editing = True supports_sampling_controls = True supports_diffusion = True model_loaded = False def __init__( self, model_id: Optional[str] = None, device: str = "cpu", hf_token: Optional[str] = None, ): self.model_id = model_id or self.model_id self.device = device self.hf_token = hf_token self.pipe = None def warmup(self) -> bool: self._ensure_loaded() return self.model_loaded def _ensure_loaded(self): if self.pipe is not None: return if self.requires_cuda and not self.device.startswith("cuda"): raise RuntimeError(f"{self.provider_label} needs a CUDA GPU. This machine is CPU-only.") self.pipe = self._load_pipeline() self.model_loaded = True def _load_pipeline(self): raise NotImplementedError def _generate( self, image: Image.Image, prompt: str, num_steps: int, text_guidance_scale: float, image_guidance_scale: float, seed: Optional[int], ) -> Image.Image: raise NotImplementedError def edit( self, image: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, mask: Optional[Image.Image] = None, reference_image: Optional[Image.Image] = None, background_image: Optional[Image.Image] = None, ) -> EditResult: self._ensure_loaded() edited = self._generate( image=image.convert("RGB"), prompt=prompt, num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, ) return EditResult( image=edited.convert("RGB"), message=f"{self.provider_label} edit complete.", provider=self.provider_id, used_fallback=False, steps=["foundation-diffusion"], ) def _load_kwargs(self) -> dict[str, Any]: kwargs: dict[str, Any] = {} if self.hf_token: kwargs["token"] = self.hf_token return kwargs class InstructPix2PixProvider(FoundationModelProvider): provider_id = "instruct-pix2pix" provider_label = "InstructPix2Pix local editor" model_id = "timbrooks/instruct-pix2pix" prompt_hint = "Instruction-based local diffusion editing." # Ensure all app modes map to InstructPix2Pix exclusively supports_generation = True supports_inpaint = True supports_style_transfer = True supports_background_ops = True supports_upscale = True def generate( self, prompt: str, width: int = 768, height: int = 768, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: blank_image = Image.new("RGB", (width, height), (128, 128, 128)) return self.edit( image=blank_image, prompt=f"generate {prompt}", num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, ) def inpaint( self, image: Image.Image, mask: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: return self.edit( image=image, prompt=prompt or "inpaint the masked region", num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, ) def style_transfer( self, image: Image.Image, reference_image: Image.Image, prompt: str = "", seed: Optional[int] = None, ) -> EditResult: p = prompt or "apply the style of the reference image" return self.edit( image=image, prompt=p, num_steps=28, text_guidance_scale=5.5, image_guidance_scale=1.2, seed=seed, ) def background_edit( self, image: Image.Image, prompt: str, background_image: Optional[Image.Image] = None, seed: Optional[int] = None, ) -> EditResult: return self.edit( image=image, prompt=prompt or "replace the background", num_steps=36, text_guidance_scale=6.0, image_guidance_scale=1.15, seed=seed, ) def upscale( self, image: Image.Image, scale: float = 2.0, prompt: str = "", seed: Optional[int] = None, ) -> EditResult: w, h = image.size upscaled = image.resize((int(w * scale), int(h * scale)), Image.LANCZOS) return self.edit( image=upscaled, prompt=prompt or "enhance details and upscale", num_steps=20, text_guidance_scale=7.5, image_guidance_scale=1.1, seed=seed, ) def _load_pipeline(self): from diffusers import StableDiffusionInstructPix2PixPipeline torch_dtype = torch.float16 if self.device.startswith("cuda") else torch.float32 pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained( self.model_id, torch_dtype=torch_dtype, **self._load_kwargs(), ) pipe.to(self.device) if hasattr(pipe, "enable_attention_slicing"): pipe.enable_attention_slicing() # Optimize scheduler for CPU if running on CPU if not self.device.startswith("cuda"): from diffusers import DPMSolverMultistepScheduler pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) return pipe def _generate( self, image: Image.Image, prompt: str, num_steps: int, text_guidance_scale: float, image_guidance_scale: float, seed: Optional[int], ) -> Image.Image: generator = None if seed is not None: generator = torch.Generator(device=self.device).manual_seed(seed) if not self.device.startswith("cuda"): # Balance speed and quality on CPU (10-14 steps to produce clear and great outputs) steps = max(10, min(num_steps, 14)) else: steps = max(20, min(num_steps, 100)) output = self.pipe( prompt=prompt, image=image.convert("RGB"), num_inference_steps=steps, guidance_scale=max(1.0, text_guidance_scale), image_guidance_scale=max(1.0, image_guidance_scale), generator=generator, ) return output.images[0].convert("RGB") class FluxKontextProvider(FoundationModelProvider): provider_id = "flux-kontext" provider_label = "FLUX Kontext local editor" model_id = "black-forest-labs/FLUX.1-Kontext-dev" prompt_hint = "Large local editing model for broad scene and object changes." requires_cuda = True def _load_pipeline(self): from diffusers import FluxKontextPipeline pipe = FluxKontextPipeline.from_pretrained( self.model_id, torch_dtype=torch.bfloat16, **self._load_kwargs(), ) pipe.to(self.device) return pipe def _generate( self, image: Image.Image, prompt: str, num_steps: int, text_guidance_scale: float, image_guidance_scale: float, seed: Optional[int], ) -> Image.Image: generator = None if seed is not None: generator = torch.Generator(device=self.device).manual_seed(seed) output = self.pipe( image=image.convert("RGB"), prompt=prompt, guidance_scale=max(1.0, min(text_guidance_scale, 4.0)), num_inference_steps=max(15, min(num_steps, 50)), generator=generator, ) return output.images[0].convert("RGB") class HybridEditProvider(BaseEditProvider): provider_id = "hybrid-local" provider_label = "Hybrid local editor" supports_broad_editing = True supports_sampling_controls = True supports_generation = True supports_inpaint = True supports_batch = True supports_style_transfer = True supports_background_ops = True supports_upscale = True prompt_hint = ( "Prompt -> parser -> image understanding -> CV planning -> optional diffusion -> " "identity preservation -> upscale -> refinement." ) def __init__(self, diffusion_backend: Optional[BaseEditProvider], cv_provider: CVDirectProvider): self.diffusion_backend = diffusion_backend self.cv_provider = cv_provider self.engine = cv_provider.engine self.orchestrator = EditingOrchestrator(engine=self.engine) self.supports_sampling_controls = diffusion_backend is not None and diffusion_backend.supports_sampling_controls self.supports_diffusion = diffusion_backend is not None and diffusion_backend.supports_diffusion if diffusion_backend is not None: self.provider_label = f"Hybrid local editor ({diffusion_backend.provider_label})" @property def model_loaded(self) -> bool: if self.diffusion_backend is None: return True return bool(getattr(self.diffusion_backend, "model_loaded", True)) def warmup(self) -> bool: cv_ready = self.cv_provider.warmup() if self.diffusion_backend is None: return cv_ready backend_ready = self.diffusion_backend.warmup() return cv_ready and backend_ready def capabilities(self) -> list[dict[str, Any]]: return self.cv_provider.capabilities() def presets(self) -> list[dict[str, Any]]: return self.cv_provider.presets() def _active_backend(self) -> BaseEditProvider: return self.diffusion_backend or self.cv_provider def edit( self, image: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, mask: Optional[Image.Image] = None, reference_image: Optional[Image.Image] = None, background_image: Optional[Image.Image] = None, ) -> EditResult: outcome = self.orchestrator.run( backend=self._active_backend(), fallback=self.cv_provider, image=image.convert("RGB"), prompt=prompt, num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, mask=mask, reference_image=reference_image, background_image=background_image, disable_diffusion=self.diffusion_backend is None, ) return EditResult( image=outcome.image.convert("RGB"), message=outcome.message, provider=self.provider_id, used_fallback=outcome.used_fallback, steps=outcome.steps, metadata=outcome.metadata, ) def generate( self, prompt: str, width: int = 768, height: int = 768, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: outcome = self.orchestrator.run_generate( backend=self._active_backend(), prompt=prompt, width=width, height=height, num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, ) return EditResult( image=outcome.image.convert("RGB"), message=outcome.message, provider=self.provider_id, used_fallback=outcome.used_fallback, steps=outcome.steps, metadata=outcome.metadata, ) def inpaint( self, image: Image.Image, mask: Image.Image, prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> EditResult: if self.diffusion_backend is not None and getattr(self.diffusion_backend, "supports_inpaint", False): try: result = self.diffusion_backend.inpaint( image=image, mask=mask, prompt=prompt or "repair the masked region naturally", num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, ) result.provider = self.provider_id return result except Exception: pass result = self.cv_provider.inpaint( image=image, mask=mask, prompt=prompt or "repair the masked region naturally", num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=seed, ) result.provider = self.provider_id result.used_fallback = True return result def batch_edit( self, images: list[Image.Image], prompt: str, num_steps: int = 50, text_guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, seed: Optional[int] = None, ) -> list[EditResult]: results: list[EditResult] = [] for idx, image in enumerate(images): current_seed = None if seed is None else seed + idx results.append( self.edit( image=image, prompt=prompt, num_steps=num_steps, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, seed=current_seed, ) ) return results def style_transfer( self, image: Image.Image, reference_image: Image.Image, prompt: str = "", seed: Optional[int] = None, ) -> EditResult: outcome = self.orchestrator.run( backend=self._active_backend(), fallback=self.cv_provider, image=image.convert("RGB"), prompt=prompt or "apply the reference style", num_steps=28, text_guidance_scale=5.5, image_guidance_scale=1.2, seed=seed, reference_image=reference_image, prepend_steps=[OperationStep(name="style_reference", params={"amount": 0.82})], disable_diffusion=self.diffusion_backend is None, ) return EditResult( image=outcome.image.convert("RGB"), message=outcome.message, provider=self.provider_id, used_fallback=outcome.used_fallback, steps=outcome.steps, metadata={"mode": "style-transfer", **outcome.metadata}, ) def background_edit( self, image: Image.Image, prompt: str, background_image: Optional[Image.Image] = None, seed: Optional[int] = None, ) -> EditResult: outcome = self.orchestrator.run( backend=self._active_backend(), fallback=self.cv_provider, image=image.convert("RGB"), prompt=prompt or "replace the background", num_steps=36, text_guidance_scale=6.0, image_guidance_scale=1.15, seed=seed, background_image=background_image, prepend_steps=[OperationStep(name="background_replace", params={"amount": 0.85})], disable_diffusion=self.diffusion_backend is None, ) return EditResult( image=outcome.image.convert("RGB"), message=outcome.message, provider=self.provider_id, used_fallback=outcome.used_fallback, steps=outcome.steps, metadata={"mode": "background", **outcome.metadata}, ) def upscale( self, image: Image.Image, scale: float = 2.0, prompt: str = "", seed: Optional[int] = None, ) -> EditResult: if self.diffusion_backend is not None and getattr(self.diffusion_backend, "supports_upscale", False): try: result = self.diffusion_backend.upscale( image=image, scale=scale, prompt=prompt, seed=seed ) result.provider = self.provider_id return result except Exception: pass result = self.cv_provider.upscale(image=image, scale=scale, prompt=prompt, seed=seed) result.provider = self.provider_id return result def create_edit_provider( provider_name: str = "auto", foundation_backend: Optional[str] = None, foundation_model_id: Optional[str] = None, checkpoint_path: Optional[str] = None, vae_checkpoint_path: Optional[str] = None, device: str = "cpu", ): provider_name = (provider_name or "auto").lower() foundation_backend = ( foundation_backend or os.getenv("IMAGE_EDIT_BACKEND") or os.getenv("LOCAL_IMAGE_EDIT_BACKEND") or "instruct-pix2pix" ).lower() foundation_model_id = foundation_model_id or os.getenv("LOCAL_IMAGE_EDIT_MODEL") hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") engine = CVEngine() cv_provider = CVDirectProvider(engine=engine) # Auto-resolve checkpoint paths on startup if not explicitly provided if provider_name in {"auto", "custom"} and not checkpoint_path: for candidate in ["checkpoints_cpu/diffusion_final.pt", "checkpoints/diffusion_final.pt", "checkpoints/model.pt"]: if os.path.exists(candidate): checkpoint_path = candidate break if provider_name in {"auto", "custom"} and not vae_checkpoint_path: for candidate in ["checkpoints_cpu/vae_final.pt", "checkpoints/vae_final.pt", "checkpoints/vae.pt"]: if os.path.exists(candidate): vae_checkpoint_path = candidate break # Replaced CPU custom auto-selection: using InstructPix2Pix by default as requested. if provider_name in {"auto", "foundation"} and foundation_backend: if foundation_backend == "flux-kontext": diffusion_backend = FluxKontextProvider(model_id=foundation_model_id, device=device, hf_token=hf_token) return HybridEditProvider(diffusion_backend=diffusion_backend, cv_provider=cv_provider) if foundation_backend == "instruct-pix2pix": diffusion_backend = InstructPix2PixProvider(model_id=foundation_model_id, device=device, hf_token=hf_token) return HybridEditProvider(diffusion_backend=diffusion_backend, cv_provider=cv_provider) raise RuntimeError(f"Unknown foundation backend: {foundation_backend}") if provider_name == "foundation" and not foundation_backend: raise RuntimeError( "Foundation provider requested but no backend was configured. " "Use IMAGE_EDIT_BACKEND=instruct-pix2pix or flux-kontext." ) if provider_name in {"auto", "custom"} and checkpoint_path and os.path.exists(checkpoint_path): from model.pipeline import EditPipeline pipeline = EditPipeline.from_checkpoint( checkpoint_path, vae_checkpoint_path=vae_checkpoint_path, device=device, ) diffusion_backend = CustomDiffusionProvider(pipeline=pipeline, fallback=cv_provider) return HybridEditProvider(diffusion_backend=diffusion_backend, cv_provider=cv_provider) if provider_name == "custom": raise RuntimeError("Custom provider requested but no valid checkpoint path was supplied.") return HybridEditProvider(diffusion_backend=None, cv_provider=cv_provider)