ArtiAgent - DefectFill/src/pipeline/defectdiffu_generator.py

"""
DefectDiffu Generator Wrapper

Replaces FLUX/RF-Solver-Edit with DefectDiffu's text-guided disentangled
architecture for manufacturing defect generation.

DefectDiffu (ECCV 2024) uses:
  - Three disentangled text prompts: c_p (product/bg), c_d (defect), c_f (fusion)
  - Double-free strategy with perturbation scales w_p and w_d
  - Automatic mask extraction from defect-block cross-attention maps
  - DiT backbone + Stable Diffusion VAE

INTEGRATION NOTE:
  This file contains the INTERFACE. You must plug in the actual DefectDiffu
  model forward pass from the official repo at the marked TODO sections.
  Repo: https://github.com/FFDD-diffusion/DefectDiffu
"""

import os
import torch
import torch.nn.functional as F
import numpy as np
from typing import Dict, Tuple, Optional, List
from dataclasses import dataclass
from PIL import Image
import warnings

# === DefectDiffu actual imports (must be in PYTHONPATH) ===
from diffusers.models import AutoencoderKL

import sys
from pathlib import Path

# Automatically locate and add the DefectDiffu engine directory to sys.path
CURRENT_DIR = Path(__file__).resolve().parent
DEFECTDIFFU_DIR = CURRENT_DIR.parent.parent / "engine" / "DefectDiffu"

if DEFECTDIFFU_DIR.exists() and str(DEFECTDIFFU_DIR) not in sys.path:
    sys.path.insert(0, str(DEFECTDIFFU_DIR))
import clip.clip as clip

from models_add_cross_concate import DiT
from diffusion import create_diffusion


# =========================================================================
# Mask binarization helpers (copied from test.py)
# =========================================================================

def rgb_to_gray(tensor):
    r, g, b = tensor[:, 0], tensor[:, 1], tensor[:, 2]
    gray = 0.299 * r + 0.587 * g + 0.114 * b
    return gray


def iterative_thresholding_batch(gray_tensor):
    gray_np = gray_tensor.detach().cpu().numpy()
    binarized = np.zeros_like(gray_np, dtype=np.uint8)

    for i in range(gray_np.shape[0]):
        img = gray_np[i]
        T = img.mean()
        prev_T = -1

        while abs(T - prev_T) > 1e-4:
            prev_T = T
            G1 = img[img >= T]
            G2 = img[img < T]
            m1 = G1.mean() if G1.size > 0 else 0
            m2 = G2.mean() if G2.size > 0 else 0
            T = (m1 + m2) / 2

        binarized[i] = (img >= T).astype(np.uint8)

    return torch.from_numpy(binarized).to(gray_tensor.device)


def binarize_tensor_iterative(x):
    gray = rgb_to_gray(x)
    binary = iterative_thresholding_batch(gray)
    return binary.unsqueeze(1)

@dataclass
class DefectDiffuConfig:
    """Configuration for DefectDiffu inference."""
    ckpt_path: str                          # Path to trained DefectDiffu checkpoint
    vae_path: str                           # Path to SD VAE (stabilityai/sd-vae-ft-mse)
    dit_model: str = "DiT-XL/2"             # DiT variant (DiT-XL/2, DiT-L/2, etc.)
    image_size: int = 512                   # Must match training resolution
    num_steps: int = 50                     # DDPM/DDIM inference steps
    cfg_scale: float = 1.0                  # Classifier-free guidance (if used)
    device: str = "cuda"
    offload: bool = False                   # CPU offload for low-VRAM GPUs
    seed: int = 42


class DefectDiffuGenerator:
    """
    Wrapper around DefectDiffu for the agentic pipeline.

    Unlike FLUX (which edits an existing image via inversion-injection),
    DefectDiffu generates a NEW image from noise conditioned on three text
    prompts. The input "clean image" is used only for planning/verification,
    not as a pixel-level source for editing.
    """

    def __init__(self, config: DefectDiffuConfig):
        self.config = config
        self.device = torch.device(config.device)
        self._models_loaded = False

        # Placeholders — populated in _load_models()
        self.dit = None
        self.vae = None
        self.text_encoder = None
        self.tokenizer = None
        self.scheduler = None

        self._load_models()

    # ------------------------------------------------------------------
    # TODO: Replace the methods below with actual DefectDiffu code
    # ------------------------------------------------------------------

    def _load_models(self):
        """Load DiT, VAE, text encoder, and scheduler."""
        print(f"[DefectDiffu] Loading checkpoint: {self.config.ckpt_path}")
        print(f"[DefectDiffu] VAE: {self.config.vae_path}")

        # 1. CLIP RN50 (must match training)
        self.model_clip, _ = clip.load('RN50', self.device)
        self.model_clip.eval()

        # 2. DiT architecture (must match train.py exactly)
        latent_size = self.config.image_size // 8
        self.dit = DiT(
            depth=28,
            hidden_size=1152,
            patch_size=2,
            num_heads=16,
            input_size=latent_size,
            num_classes=1000
        ).to(self.device)

        print(f"[DefectDiffu] Loading DiT weights from: {self.config.ckpt_path}")
        checkpoint = torch.load(self.config.ckpt_path, map_location=self.device)
        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
            self.dit.load_state_dict(checkpoint['model_state_dict'])
        else:
            self.dit.load_state_dict(checkpoint)
        self.dit.eval()

        # 3. Stable Diffusion VAE
        self.vae = AutoencoderKL.from_pretrained(self.config.vae_path).to(self.device)
        self.vae.eval()

        # 4. Diffusion sampler (respacing = num_steps)
        self.diffusion = create_diffusion(timestep_respacing=str(self.config.num_steps))

        self._models_loaded = True
        print("[DefectDiffu] All models loaded successfully.")

    def _encode_text(self, prompt: str) -> torch.Tensor:
        """Encode a text prompt into CLIP RN50 text embeddings."""
        with torch.no_grad():
            tokens = clip.tokenize([prompt]).to(self.device)
            emb = self.model_clip.encode_text(tokens)
            emb = emb / emb.norm(dim=-1, keepdim=True)
            emb = emb.float()
        return emb

    def _extract_mask_from_attention(
        self,
        mask_latent: torch.Tensor
    ) -> np.ndarray:
        """
        Decode mask latent through VAE and binarize using iterative thresholding.
        Matches test.py post-processing.
        """
        with torch.no_grad():
            mask_decoded = self.vae.decode(mask_latent / 0.18215).sample  # [1, 3, H, W]

        # Binarize with iterative thresholding (Otsu-like)
        mask_binary = binarize_tensor_iterative(mask_decoded)  # [1, 1, H, W]
        mask_bool = mask_binary[0, 0].cpu().numpy() > 0

        return mask_bool

    def _denoise_with_double_free(
        self,
        z: torch.Tensor,
        emb_p: torch.Tensor,
        emb_d: torch.Tensor,
        emb_f: torch.Tensor,
        emb_good: torch.Tensor,
        emb_null_good: torch.Tensor,
        w_d: float,
        w_p: float
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Run DefectDiffu inference via p_sample_loop with dual-branch CFG.
        Matches test.py exactly.

        Returns:
            (img_latent, mask_latent) in VAE latent space, shape [1, 4, H, W]
        """
        # Build paired conditioning: defect_class vs good_class
        y_defect_class = [emb_d, emb_p, emb_f]
        y_good_class = [emb_good, emb_p, emb_null_good]
        y = [y_defect_class, y_good_class]

        # Duplicate latent for CFG (concatenated batch)
        z_cfg = torch.cat([z, z], dim=0)

        model_kwargs = dict(y=y, cfg_scale=float(w_d))

        with torch.no_grad():
            samples, cross = self.diffusion.p_sample_loop(
                self.dit.forward_with_cfg_2,
                z_cfg.shape,
                z_cfg,
                clip_denoised=False,
                model_kwargs=model_kwargs,
                progress=False,
                device=self.device
            )

        # Unchunk: first half is the defect-conditioned output
        img_latent, _ = samples.chunk(2, dim=0)
        mask_latent, _ = cross.chunk(2, dim=0)

        return img_latent, mask_latent

    # ------------------------------------------------------------------
    # Public API — used by the orchestrator
    # ------------------------------------------------------------------

    @torch.no_grad()
    def generate(
        self,
        c_p: str,
        c_d: str,
        c_f: str,
        w_d: float = 1.0,
        w_p: float = 1.0,
        seed: Optional[int] = None
    ) -> Tuple[Image.Image, np.ndarray]:
        """
        Generate a synthetic defect image and its binary mask.

        Args:
            c_p: Background/product prompt  (e.g. "A photo of metal nut")
            c_d: Defect prompt                (e.g. "A photo of scratch")
            c_f: Fusion prompt                (e.g. "A photo of metal nut with scratch")
            w_d: Defect strength perturbation scale (0.0 = no defect, 2.0 = severe)
            w_p: Product consistency scale    (usually 1.0, increase for stronger product fidelity)
            seed: Random seed

        Returns:
            (pil_image, binary_mask) where mask is bool array [H, W]
        """
        if seed is None:
            seed = self.config.seed
        torch.manual_seed(seed)
        np.random.seed(seed)

        print(f"[DefectDiffu] Generating: w_d={w_d}, w_p={w_p}")
        print(f"[DefectDiffu]   c_p: {c_p}")
        print(f"[DefectDiffu]   c_d: {c_d}")
        print(f"[DefectDiffu]   c_f: {c_f}")

                # 1. Parse product name from c_p for the null-good prompt
        product_name = c_p.replace("A photo of ", "").strip()

        # 2. Encode all five text conditions (must match training format)
        emb_d = self._encode_text(c_d)                       # "a photo of scratch"
        emb_p = self._encode_text(c_p)                       # "a photo of vcsel"
        emb_f = self._encode_text(c_f)                       # "a photo of scratch vcsel"
        emb_good = self._encode_text("a photo of good")
        emb_null_good = self._encode_text(f"a photo of good {product_name}")

        # 3. Initialize latent noise
        latent_h = self.config.image_size // 8
        latent_w = self.config.image_size // 8
        z = torch.randn(1, 4, latent_h, latent_w, device=self.device)

        # 4. DefectDiffu double-free denoising
        img_latent, mask_latent = self._denoise_with_double_free(
            z, emb_p, emb_d, emb_f, emb_good, emb_null_good, w_d, w_p
        )

        # 5. Decode image latent → RGB
        with torch.no_grad():
            img_tensor = self.vae.decode(img_latent / 0.18215).sample  # [1, 3, H, W]
        img_tensor = (img_tensor + 1) / 2  # [-1, 1] → [0, 1]
        img_tensor = img_tensor.clamp(0, 1)

        img_np = (img_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
        pil_image = Image.fromarray(img_np)

        # 6. Extract defect mask from mask latent
        binary_mask = self._extract_mask_from_attention(mask_latent)

        print(f"[DefectDiffu] Generation complete. Mask coverage: {binary_mask.mean():.3f}")
        return pil_image, binary_mask

    @torch.no_grad()
    def generate_from_plan(
        self,
        product_description: str,
        defect_description: str,
        severity: str = "medium",
        w_d: Optional[float] = None,
        w_p: float = 1.0,
        seed: Optional[int] = None
    ) -> Tuple[Image.Image, np.ndarray, Dict]:
        """
        Convenience wrapper that builds the three DefectDiffu prompts from
        product/defect descriptions and maps severity to w_d.
        """
        # Map severity to defect strength
        severity_to_wd = {"low": 0.6, "minor": 0.6,
                          "medium": 1.0, "moderate": 1.0,
                          "high": 1.5, "severe": 1.5}
        if w_d is None:
            w_d = severity_to_wd.get(severity.lower(), 1.0)

        c_p = f"A photo of {product_description}"
        c_d = f"A photo of {defect_description}"
        c_f = f"A photo of {product_description} with {defect_description}"

        img, mask = self.generate(c_p, c_d, c_f, w_d=w_d, w_p=w_p, seed=seed)

        meta = {
            "c_p": c_p, "c_d": c_d, "c_f": c_f,
            "w_d": w_d, "w_p": w_p, "seed": seed or self.config.seed
        }
        return img, mask, meta

    def unload_models(self):
        """Free GPU memory."""
        self.dit = None
        self.vae = None
        self.model_clip = None
        self.diffusion = None
        if self.device.type == "cuda":
            torch.cuda.empty_cache()
        print("[DefectDiffu] Models unloaded.")


def get_defectdiffu_generator(
    ckpt_path: str,
    vae_path: str,
    device: str = "cuda",
    **kwargs
) -> DefectDiffuGenerator:
    """Factory function for easy instantiation."""
    config = DefectDiffuConfig(ckpt_path=ckpt_path, vae_path=vae_path,
                                device=device, **kwargs)
    return DefectDiffuGenerator(config)