"""Generative API clients for double-exposure separation.""" from __future__ import annotations import os from dataclasses import dataclass from typing import List, Optional, Tuple import numpy as np from PIL import Image from app.preprocessing import to_pil DEFAULT_SEPARATION_MODEL = os.environ.get( "REPLICATE_SEPARATION_MODEL", "black-forest-labs/flux-dev", ) DEMO_PERCENTILES = (35.0, 45.0, 50.0, 55.0, 65.0) REPLICATE_STRENGTHS = (0.55, 0.65, 0.75) @dataclass class SeparationResult: """Output from a separation attempt.""" image_a: np.ndarray # float RGB [0, 1] image_b: np.ndarray method: str message: str candidate_id: str = "default" scan_analysis: Optional["ScanAnalysis"] = None # WP-5.1 Fix 7: carry VLM analysis to avoid re-call diagnostics: Optional[dict] = None # WP-6: structured per-source numbers (e.g. DIP init/best loss) def _luminance(rgb: np.ndarray) -> np.ndarray: return ( 0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2] ) def _apply_hard_mask(base: np.ndarray, mask: np.ndarray, fill: float = 0.35) -> np.ndarray: gray = np.mean(base, axis=-1, keepdims=True) out = base * mask[..., np.newaxis] + gray * (~mask[..., np.newaxis]) * fill return np.clip(out, 0.0, 1.0).astype(np.float32) def _demo_separation_percentile( positive_rgb: np.ndarray, percentile: float, candidate_id: str, soft: bool = False, ) -> SeparationResult: """Split by luminance percentile; optional soft Gaussian boundary.""" lum = _luminance(positive_rgb) threshold = float(np.percentile(lum, percentile)) if soft: sigma = max(threshold * 0.15, 0.02) weight_a = np.exp(-((lum - threshold) ** 2) / (2 * sigma**2)) weight_a = np.where(lum <= threshold, 1.0, weight_a) weight_b = 1.0 - weight_a image_a = np.clip(positive_rgb * weight_a[..., np.newaxis], 0.0, 1.0) image_b = np.clip(positive_rgb * weight_b[..., np.newaxis], 0.0, 1.0) strategy = f"soft_p{int(percentile)}" else: mask_a = lum <= threshold image_a = _apply_hard_mask(positive_rgb, mask_a) image_b = _apply_hard_mask(positive_rgb, ~mask_a) strategy = f"hard_p{int(percentile)}" return SeparationResult( image_a=image_a.astype(np.float32), image_b=image_b.astype(np.float32), method=f"demo_{strategy}", message="Demo heuristic separation (no API key).", candidate_id=candidate_id, ) def _demo_separation_spatial(positive_rgb: np.ndarray, axis: str) -> SeparationResult: """Split along horizontal or vertical midline with feathered blend.""" h, w = positive_rgb.shape[:2] if axis == "horizontal": coord = np.linspace(0, 1, h)[:, np.newaxis] coord = np.broadcast_to(coord, (h, w)) cid = "spatial_h" else: coord = np.linspace(0, 1, w)[np.newaxis, :] coord = np.broadcast_to(coord, (h, w)) cid = "spatial_v" weight_a = np.clip(1.0 - np.abs(coord - 0.5) * 4.0, 0.0, 1.0) weight_b = 1.0 - weight_a image_a = np.clip(positive_rgb * weight_a[..., np.newaxis], 0.0, 1.0) image_b = np.clip(positive_rgb * weight_b[..., np.newaxis], 0.0, 1.0) return SeparationResult( image_a=image_a.astype(np.float32), image_b=image_b.astype(np.float32), method=f"demo_{cid}", message="Demo heuristic separation (no API key).", candidate_id=cid, ) def generate_demo_candidates( positive_rgb: np.ndarray, num_candidates: int = 3, ) -> List[SeparationResult]: """Generate diverse demo candidates without API access.""" pool: List[SeparationResult] = [] for p in DEMO_PERCENTILES: pool.append( _demo_separation_percentile( positive_rgb, p, candidate_id=f"hard_p{int(p)}", soft=False ) ) pool.append( _demo_separation_percentile( positive_rgb, 50.0, candidate_id="soft_p50", soft=True ) ) pool.extend([ _demo_separation_spatial(positive_rgb, "horizontal"), _demo_separation_spatial(positive_rgb, "vertical"), ]) return pool[: max(1, min(num_candidates, len(pool)))] def _replicate_separation( positive_rgb: np.ndarray, model: str = DEFAULT_SEPARATION_MODEL, prompt_strength: float = 0.65, candidate_id: str = "replicate_0", ) -> SeparationResult: """Call Replicate for generative separation via image-to-image prompting.""" import replicate pil = to_pil(positive_rgb) prompt = ( "Separate this double-exposed photograph into the two distinct original " "scenes. Recover clear, photorealistic details from each exposure." ) output = replicate.run( model, input={ "prompt": prompt, "image": pil, "prompt_strength": prompt_strength, "num_inference_steps": 28, "guidance": 3.5, }, ) if isinstance(output, list): url = str(output[0]) else: url = str(output) from io import BytesIO import urllib.request with urllib.request.urlopen(url) as resp: gen_img = Image.open(BytesIO(resp.read())).convert("RGB") gen_rgb = np.asarray(gen_img, dtype=np.float32) / 255.0 if gen_rgb.shape[:2] != positive_rgb.shape[:2]: gen_pil = gen_img.resize( (positive_rgb.shape[1], positive_rgb.shape[0]), Image.Resampling.LANCZOS, ) gen_rgb = np.asarray(gen_pil, dtype=np.float32) / 255.0 lum_gen = _luminance(gen_rgb) lum_orig = _luminance(positive_rgb) lum_b = np.clip(lum_orig - lum_gen * 0.5, 0.0, 1.0) scale = lum_b[..., np.newaxis] / np.clip(lum_orig[..., np.newaxis], 1e-4, 1.0) image_b = np.clip(positive_rgb * scale, 0.0, 1.0).astype(np.float32) return SeparationResult( image_a=gen_rgb.astype(np.float32), image_b=image_b, method=f"replicate:{model}@s{prompt_strength:.2f}", message=f"Generative separation via {model} (strength={prompt_strength:.2f})", candidate_id=candidate_id, ) def generate_replicate_candidates( positive_rgb: np.ndarray, num_candidates: int = 3, model: str = DEFAULT_SEPARATION_MODEL, ) -> List[SeparationResult]: """Generate multiple Replicate separations with varied prompt strength.""" strengths = REPLICATE_STRENGTHS[: max(1, num_candidates)] results: List[SeparationResult] = [] for i, strength in enumerate(strengths): try: results.append( _replicate_separation( positive_rgb, model=model, prompt_strength=strength, candidate_id=f"replicate_s{int(strength * 100)}", ) ) except Exception as exc: results.append( SeparationResult( image_a=positive_rgb.copy(), image_b=positive_rgb.copy(), method="replicate_error", message=f"Candidate failed: {exc}", candidate_id=f"replicate_fail_{i}", ) ) return results def _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, img2img, vlm, method): """WP-5.1 Fix 8 helper to dedup append logic.""" if h_total is None or confidence_mask is None: return try: from app.demix import analyze_scan, residual_demix, DemixConfig analysis = analyze_scan(positive_rgb, vlm=vlm) cfg = DemixConfig(iterations=2, strength=0.55, use_instruct_edit=("instruct" in method)) d = residual_demix(positive_rgb, h_total, confidence_mask, img2img, analysis, cfg, method=method) if "instruct" in method: d.candidate_id = f"demix_instruct_k{analysis.k_judgment:.1f}_i{cfg.iterations}" candidates.append(d) except Exception: pass def _append_deep_prior( candidates, positive_rgb, density, log_exposure, confidence_mask, film_curve, dip_policy=None, ): """WP-6 helper: append the Double-DIP candidate (soft-fail, lazy import). Skipped silently when density/log_exposure are unavailable (same rule as demix: never optimize against the legacy circular objective). WP-13.1: dip_policy defaults to DEFAULT_POLICY (bench-safe); app passes APP_POLICY. """ if density is None or log_exposure is None or confidence_mask is None: return try: from baselines.double_dip import double_dip_separate, DoubleDIPConfig from scoring_policy import DEFAULT_POLICY if film_curve is None: from film_physics import get_film_curve film_curve = get_film_curve("Generic") cfg = DoubleDIPConfig(policy=dip_policy or DEFAULT_POLICY) res = double_dip_separate( positive_rgb, log_exposure, density, confidence_mask, film_curve, config=cfg ) if res is not None: candidates.append(res) except Exception: pass def generate_candidates( positive_rgb: np.ndarray, num_candidates: int = 3, api_token: Optional[str] = None, model: str = DEFAULT_SEPARATION_MODEL, h_total: Optional[np.ndarray] = None, confidence_mask: Optional[np.ndarray] = None, density: Optional[np.ndarray] = None, log_exposure=None, include_deep_prior: bool = False, film_curve=None, dip_policy=None, ) -> Tuple[List[SeparationResult], str]: """ Generate multiple separation candidates for ranking. h_total / confidence_mask (WP-5): when provided, demix source is registered (stub in demo; replicate+ instruct variant when token). density / log_exposure / include_deep_prior / film_curve (WP-6): when the flag is set and the density path is live, the Double-DIP source is registered. film_curve should be the user's selected stock so DIP optimizes the right physics. Returns: Tuple of (candidate list, mode description string). """ num_candidates = max(1, min(int(num_candidates), 5)) token = api_token or os.environ.get("REPLICATE_API_TOKEN", "").strip() if not token: candidates = generate_demo_candidates(positive_rgb, num_candidates) if h_total is not None and confidence_mask is not None: # WP-5.1 Fix 4+8: demo uses helper, vlm=None from app.demix import stub_cleanup _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, stub_cleanup, None, "demix_stub") if include_deep_prior: _append_deep_prior( candidates, positive_rgb, density, log_exposure, confidence_mask, film_curve, dip_policy=dip_policy, ) return candidates, "demo" os.environ["REPLICATE_API_TOKEN"] = token try: candidates = generate_replicate_candidates(positive_rgb, num_candidates, model) if all(c.method == "replicate_error" for c in candidates): raise RuntimeError("All Replicate candidates failed") # WP-5.1 Fix 3+4+8: append via helper; gate instruct if h_total is not None and confidence_mask is not None: from app.demix import replicate_img2img as _rep_img2img vlm = (lambda: __import__("app.demix", fromlist=["anthropic_vlm"]).anthropic_vlm if os.environ.get("ANTHROPIC_API_KEY") else None)() _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, _rep_img2img, vlm, "demix_replicate") if num_candidates >= 3: _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, _rep_img2img, vlm, "demix_instruct") if include_deep_prior: _append_deep_prior( candidates, positive_rgb, density, log_exposure, confidence_mask, film_curve, dip_policy=dip_policy, ) return candidates, "replicate" except Exception as exc: candidates = generate_demo_candidates(positive_rgb, num_candidates) for c in candidates: c.message = f"Replicate unavailable ({exc}). Using demo candidate." if h_total is not None and confidence_mask is not None: # WP-5.1 Fix 4+8: fallback demo via helper, vlm=None from app.demix import stub_cleanup _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, stub_cleanup, None, "demix_stub") if include_deep_prior: _append_deep_prior( candidates, positive_rgb, density, log_exposure, confidence_mask, film_curve, dip_policy=dip_policy, ) return candidates, "demo_fallback" def result_to_pil_pair(result: SeparationResult) -> Tuple[Image.Image, Image.Image]: """Convert separation arrays to PIL images for display.""" return to_pil(result.image_a), to_pil(result.image_b)