pipeline.py

#!/usr/bin/env python3
"""
Standalone CV Synthetic Data Engine

This module implements an API-first synthetic data pipeline for few-shot object
conditioning, prompt-driven synthetic frame generation, Grounding DINO zero-shot
annotation, and dataset compilation for YOLO or COCO-style training workflows.

The script is intentionally self-contained. It can be used directly from a
GPU host, wrapped by Modal serverless functions, or launched inside RunPod.
"""

from __future__ import annotations

import argparse
import gc
import json
import logging
import math
import os
import random
import shutil
import sys
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union

import numpy as np
from PIL import Image, ImageEnhance, ImageFilter, ImageOps

try:
    import cv2
except Exception as exc:  # pragma: no cover - dependency guard
    cv2 = None
    _CV2_IMPORT_ERROR = exc
else:
    _CV2_IMPORT_ERROR = None

try:
    import torch
    import torch.nn.functional as F
    from torch.utils.data import DataLoader, Dataset
except Exception as exc:  # pragma: no cover - dependency guard
    torch = None
    F = None
    DataLoader = object
    Dataset = object
    _TORCH_IMPORT_ERROR = exc
else:
    _TORCH_IMPORT_ERROR = None

try:
    from diffusers import DDPMScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline
    from peft.utils import get_peft_model_state_dict
except Exception as exc:  # pragma: no cover - dependency guard
    DDPMScheduler = None
    StableDiffusionPipeline = None
    StableDiffusionXLPipeline = None
    get_peft_model_state_dict = None
    _DIFFUSERS_IMPORT_ERROR = exc
else:
    _DIFFUSERS_IMPORT_ERROR = None

try:
    from peft import LoraConfig
except Exception as exc:  # pragma: no cover - dependency guard
    LoraConfig = None
    _PEFT_IMPORT_ERROR = exc
else:
    _PEFT_IMPORT_ERROR = None

try:
    from transformers import AutoModelForZeroShotObjectDetection, AutoProcessor
except Exception as exc:  # pragma: no cover - dependency guard
    AutoModelForZeroShotObjectDetection = None
    AutoProcessor = None
    _TRANSFORMERS_IMPORT_ERROR = exc
else:
    _TRANSFORMERS_IMPORT_ERROR = None

try:
    import modal
except Exception:  # pragma: no cover - optional platform integration
    modal = None

LOGGER = logging.getLogger("synthetic_cv_pipeline")

LabelFormat = Literal["yolo", "coco"]
ImageLike = Union[str, Path, Image.Image, np.ndarray]


@dataclass
class TrainingConfig:
    """Configuration for high-velocity few-shot LoRA optimization."""

    pretrained_model: str = "runwayml/stable-diffusion-v1-5"
    output_dir: str = "conditioned_lora"
    instance_token: str = "sksobj"
    resolution: int = 512
    train_steps: int = 180
    validation_interval: int = 20
    patience: int = 4
    min_delta: float = 0.0025
    learning_rate: float = 1e-4
    batch_size: int = 1
    gradient_accumulation_steps: int = 1
    rank: int = 8
    seed: int = 1337
    mixed_precision: Literal["fp16", "bf16", "no"] = "fp16"
    negative_prompt: str = "low quality, blurry, warped object, extra object, text, watermark"
    num_validation_images: int = 1
    train_text_encoder_lora: bool = False
    max_grad_norm: float = 1.0


@dataclass
class SynthesisConfig:
    """Configuration for conditioned batch synthesis."""

    prompt: str
    target_object: str
    output_dir: str = "output_batch"
    lora_dir: Optional[str] = "conditioned_lora"
    pretrained_model: str = "runwayml/stable-diffusion-v1-5"
    num_images: int = 24
    batch_size: int = 1
    width: int = 512
    height: int = 512
    inference_steps: int = 32
    guidance_scale: float = 7.0
    lora_scale: float = 0.85
    seed: int = 1337
    label_format: LabelFormat = "yolo"
    class_id: int = 0
    category_id: int = 1
    detector_model: str = "IDEA-Research/grounding-dino-tiny"
    detection_threshold: float = 0.35
    text_threshold: float = 0.25
    debug_preview_count: int = 5
    negative_prompt: str = "low quality, blurry, duplicate object, malformed, noisy, text, watermark"


@dataclass
class DetectionRecord:
    """Normalized internal representation of one detector output."""

    image_id: int
    label: str
    score: float
    box_xyxy: Tuple[float, float, float, float]
    width: int
    height: int

    def clipped(self) -> "DetectionRecord":
        xmin, ymin, xmax, ymax = self.box_xyxy
        xmin = float(max(0.0, min(xmin, self.width - 1)))
        ymin = float(max(0.0, min(ymin, self.height - 1)))
        xmax = float(max(0.0, min(xmax, self.width - 1)))
        ymax = float(max(0.0, min(ymax, self.height - 1)))
        if xmax < xmin:
            xmin, xmax = xmax, xmin
        if ymax < ymin:
            ymin, ymax = ymax, ymin
        return DetectionRecord(
            image_id=self.image_id,
            label=self.label,
            score=self.score,
            box_xyxy=(xmin, ymin, xmax, ymax),
            width=self.width,
            height=self.height,
        )


def configure_logging(level: str = "INFO") -> None:
    numeric_level = getattr(logging, level.upper(), logging.INFO)
    logging.basicConfig(
        level=numeric_level,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def require_dependency(name: str, import_error: Optional[BaseException]) -> None:
    if import_error is not None:
        raise RuntimeError(
            f"Required dependency '{name}' could not be imported. Install the expected GPU stack "
            f"before running this pipeline. Original error: {import_error}"
        ) from import_error


def resolve_device() -> str:
    require_dependency("torch", _TORCH_IMPORT_ERROR)
    if torch.cuda.is_available():
        return "cuda"
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def clear_vram() -> None:
    gc.collect()
    if torch is not None and torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()


def load_source_images(images: Sequence[ImageLike]) -> List[Image.Image]:
    loaded: List[Image.Image] = []
    for idx, item in enumerate(images):
        try:
            if isinstance(item, Image.Image):
                image = item.convert("RGB")
            elif isinstance(item, np.ndarray):
                arr = item
                if arr.ndim == 2:
                    arr = np.stack([arr] * 3, axis=-1)
                if arr.shape[-1] == 4:
                    arr = arr[..., :3]
                image = Image.fromarray(arr.astype(np.uint8)).convert("RGB")
            else:
                path = Path(item).expanduser().resolve()
                if not path.exists():
                    raise FileNotFoundError(f"Source image does not exist: {path}")
                image = Image.open(path).convert("RGB")
            loaded.append(image)
        except Exception as exc:
            raise ValueError(f"Failed to load source image at index {idx}: {exc}") from exc
    if not loaded:
        raise ValueError("At least one source target object image is required.")
    return loaded


def load_image_paths(directory: Union[str, Path]) -> List[Path]:
    root = Path(directory).expanduser().resolve()
    if not root.exists() or not root.is_dir():
        raise FileNotFoundError(f"Source image directory not found: {root}")
    allowed = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
    paths = sorted(path for path in root.iterdir() if path.suffix.lower() in allowed)
    if not paths:
        raise FileNotFoundError(f"No supported images found in {root}")
    return paths


def pil_to_tensor(image: Image.Image, resolution: int) -> "torch.Tensor":
    image = ImageOps.exif_transpose(image).convert("RGB")
    image = ImageOps.fit(image, (resolution, resolution), method=Image.Resampling.LANCZOS)
    arr = np.asarray(image).astype(np.float32) / 127.5 - 1.0
    tensor = torch.from_numpy(arr).permute(2, 0, 1)
    return tensor


class FewShotImageDataset(Dataset):
    """A deterministic plus stochastic few-shot image dataset for LoRA training."""

    def __init__(self, images: Sequence[Image.Image], prompt: str, resolution: int, length: int = 2048) -> None:
        self.images = list(images)
        self.prompt = prompt
        self.resolution = resolution
        self.length = max(length, len(self.images))

    def __len__(self) -> int:
        return self.length

    def _augment(self, image: Image.Image) -> Image.Image:
        image = ImageOps.exif_transpose(image).convert("RGB")
        if random.random() < 0.5:
            image = ImageOps.mirror(image)
        brightness = random.uniform(0.82, 1.18)
        contrast = random.uniform(0.85, 1.15)
        saturation = random.uniform(0.82, 1.2)
        image = ImageEnhance.Brightness(image).enhance(brightness)
        image = ImageEnhance.Contrast(image).enhance(contrast)
        image = ImageEnhance.Color(image).enhance(saturation)
        angle = random.uniform(-8, 8)
        image = image.rotate(angle, resample=Image.Resampling.BICUBIC, expand=False, fillcolor=(127, 127, 127))
        return image

    def __getitem__(self, index: int) -> Dict[str, Any]:
        image = self.images[index % len(self.images)]
        return {"pixel_values": pil_to_tensor(self._augment(image), self.resolution), "prompt": self.prompt}


def compute_structural_loss(candidate: Image.Image, references: Sequence[Image.Image], resolution: int) -> float:
    """
    Compute a lightweight structural validation loss from edge maps and luminance.

    The validation signal intentionally emphasizes shape retention instead of exact
    background fidelity. This lets early stopping preserve target object features
    while avoiding overfitting to source lighting and environment.
    """
    require_dependency("opencv-python", _CV2_IMPORT_ERROR)
    candidate_gray = np.asarray(ImageOps.fit(candidate.convert("L"), (resolution, resolution), Image.Resampling.LANCZOS))
    candidate_edges = cv2.Canny(candidate_gray, 80, 160).astype(np.float32) / 255.0
    candidate_luma = candidate_gray.astype(np.float32) / 255.0
    best_loss = float("inf")
    for ref in references:
        ref_gray = np.asarray(ImageOps.fit(ref.convert("L"), (resolution, resolution), Image.Resampling.LANCZOS))
        ref_edges = cv2.Canny(ref_gray, 80, 160).astype(np.float32) / 255.0
        ref_luma = ref_gray.astype(np.float32) / 255.0
        edge_loss = float(np.mean((candidate_edges - ref_edges) ** 2))
        luma_loss = float(np.mean((candidate_luma - ref_luma) ** 2))
        best_loss = min(best_loss, 0.75 * edge_loss + 0.25 * luma_loss)
    return best_loss


def create_prompt_variation(base_prompt: str, target_object: str, instance_token: str, index: int) -> str:
    camera_angles = [
        "front three-quarter view",
        "low angle macro shot",
        "high angle inspection view",
        "side profile perspective",
        "telephoto compressed perspective",
        "wide-angle close pass",
    ]
    illumination = [
        "soft diffuse overcast lighting",
        "hard rim light with long shadows",
        "cool fluorescent industrial illumination",
        "warm golden hour side light",
        "dramatic backlight and controlled reflections",
        "mixed practical lights with subtle glare",
    ]
    materials = [
        "mild specular reflections",
        "matte surface response",
        "gloss highlights on nearby surfaces",
        "wet floor reflections",
        "dusty atmospheric scattering",
        "clean studio-grade clarity",
    ]
    distance = [
        "object occupying 12 percent of the frame",
        "object occupying 25 percent of the frame",
        "object occupying 40 percent of the frame",
        "object in the near foreground",
        "object at medium distance",
        "object partially framed by environmental structures",
    ]
    occlusion = [
        "unoccluded target object",
        "subtle foreground occlusion at one edge",
        "partial shadow crossing the target object",
        "thin cable-like occluder in foreground",
        "minor motion blur in the environment only",
        "clean silhouette with no occlusion",
    ]
    rng = random.Random(index * 7919 + len(base_prompt))
    descriptors = [
        rng.choice(camera_angles),
        rng.choice(illumination),
        rng.choice(materials),
        rng.choice(distance),
        rng.choice(occlusion),
    ]
    descriptor_text = ", ".join(descriptors)
    return f"a photo of {instance_token} {target_object} in {base_prompt}, {descriptor_text}, realistic, high detail"


def encode_prompt_for_training(pipe: Any, prompt: Union[str, List[str]], device: str) -> "torch.Tensor":
    tokens = pipe.tokenizer(
        prompt,
        padding="max_length",
        max_length=pipe.tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    )
    input_ids = tokens.input_ids.to(device)
    return pipe.text_encoder(input_ids)[0]


class FewShotLoRATrainer:
    """Few-shot LoRA trainer with early stopping on validation structural loss."""

    def __init__(self, config: TrainingConfig) -> None:
        self.config = config
        self.device = resolve_device()
        self.dtype = self._resolve_dtype()
        self.pipe: Optional[Any] = None

    def _resolve_dtype(self) -> "torch.dtype":
        if torch is None:
            raise RuntimeError("PyTorch is required for training.")
        if self.config.mixed_precision == "bf16" and torch.cuda.is_available() and torch.cuda.is_bf16_supported():
            return torch.bfloat16
        if self.config.mixed_precision == "fp16" and self.device == "cuda":
            return torch.float32
        return torch.float32

    def _load_pipeline(self) -> Any:
        require_dependency("diffusers", _DIFFUSERS_IMPORT_ERROR)
        require_dependency("peft", _PEFT_IMPORT_ERROR)
        LOGGER.info("Loading diffusion pipeline for LoRA training: %s", self.config.pretrained_model)
        try:
            pipe = StableDiffusionPipeline.from_pretrained(
                self.config.pretrained_model,
                torch_dtype=self.dtype,
                safety_checker=None,
                requires_safety_checker=False,
            )
            pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
            pipe.to(self.device)
            pipe.vae.requires_grad_(False)
            pipe.text_encoder.requires_grad_(False)
            pipe.unet.requires_grad_(False)
            lora_config = LoraConfig(
                r=self.config.rank,
                lora_alpha=self.config.rank,
                init_lora_weights="gaussian",
                target_modules=["to_k", "to_q", "to_v", "to_out.0"],
            )
            pipe.unet.add_adapter(lora_config)
            if self.config.train_text_encoder_lora:
                text_lora_config = LoraConfig(
                    r=self.config.rank,
                    lora_alpha=self.config.rank,
                    init_lora_weights="gaussian",
                    target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
                )
                pipe.text_encoder.add_adapter(text_lora_config)
            if hasattr(pipe, "enable_xformers_memory_efficient_attention"):
                try:
                    pipe.enable_xformers_memory_efficient_attention()
                    LOGGER.info("Enabled xFormers memory-efficient attention.")
                except Exception as exc:
                    LOGGER.warning("Could not enable xFormers attention: %s", exc)
            if hasattr(pipe, "enable_attention_slicing"):
                pipe.enable_attention_slicing()
            self.pipe = pipe
            return pipe
        except torch.cuda.OutOfMemoryError as exc:
            clear_vram()
            raise RuntimeError("CUDA VRAM exhausted while loading the diffusion training pipeline.") from exc
        except Exception as exc:
            clear_vram()
            raise RuntimeError(f"Failed to load diffusion training pipeline: {exc}") from exc

    def train(self, source_images: Sequence[ImageLike], target_object: str) -> Path:
        images = load_source_images(source_images)
        random.seed(self.config.seed)
        np.random.seed(self.config.seed)
        torch.manual_seed(self.config.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(self.config.seed)

        output_dir = Path(self.config.output_dir).expanduser().resolve()
        output_dir.mkdir(parents=True, exist_ok=True)
        prompt = f"a photo of {self.config.instance_token} {target_object}"
        dataset = FewShotImageDataset(images, prompt=prompt, resolution=self.config.resolution)
        loader = DataLoader(dataset, batch_size=self.config.batch_size, shuffle=True, num_workers=0)
        iterator = iter(loader)
        pipe = self._load_pipeline()

        trainable_params = [p for p in pipe.unet.parameters() if p.requires_grad]
        if self.config.train_text_encoder_lora:
            trainable_params += [p for p in pipe.text_encoder.parameters() if p.requires_grad]
        if not trainable_params:
            raise RuntimeError("No trainable LoRA parameters were registered. Check PEFT and Diffusers versions.")
        optimizer = torch.optim.AdamW(trainable_params, lr=self.config.learning_rate, betas=(0.9, 0.999), weight_decay=0.01)
        scaler_enabled = self.dtype == torch.float32 and self.device == "cuda"
        scaler = torch.cuda.amp.GradScaler(enabled=False)

        best_structural_loss = float("inf")
        patience_counter = 0
        global_step = 0
        running_loss = 0.0
        start = time.time()
        pipe.unet.train()
        if self.config.train_text_encoder_lora:
            pipe.text_encoder.train()
        LOGGER.info("Starting LoRA optimization with %d source images for target '%s'.", len(images), target_object)

        while global_step < self.config.train_steps:
            try:
                batch = next(iterator)
            except StopIteration:
                iterator = iter(loader)
                batch = next(iterator)
            pixel_values = batch["pixel_values"].to(device=self.device, dtype=self.dtype)
            prompts = list(batch["prompt"])

            with torch.no_grad():
                latents = pipe.vae.encode(pixel_values).latent_dist.sample()
                latents = latents * pipe.vae.config.scaling_factor
                noise = torch.randn_like(latents)
                timesteps = torch.randint(
                    0,
                    pipe.scheduler.config.num_train_timesteps,
                    (latents.shape[0],),
                    device=self.device,
                    dtype=torch.long,
                )
                noisy_latents = pipe.scheduler.add_noise(latents, noise, timesteps)
                encoder_hidden_states = encode_prompt_for_training(pipe, prompts, self.device).to(dtype=self.dtype)

            try:
                with torch.amp.autocast('cuda', dtype=torch.float32):
                    model_pred = pipe.unet(noisy_latents, timesteps, encoder_hidden_states).sample
                    target = noise
                    if getattr(pipe.scheduler.config, "prediction_type", None) == "v_prediction":
                        target = pipe.scheduler.get_velocity(latents, noise, timesteps)
                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
                    loss = loss / self.config.gradient_accumulation_steps
                scaler.scale(loss).backward()
            except torch.cuda.OutOfMemoryError as exc:
                clear_vram()
                raise RuntimeError("CUDA VRAM exhausted during LoRA optimization. Reduce resolution, rank, or batch size.") from exc

            running_loss += float(loss.detach().cpu().item())
            if (global_step + 1) % self.config.gradient_accumulation_steps == 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(trainable_params, self.config.max_grad_norm)
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)

            global_step += 1
            if global_step % max(1, self.config.validation_interval) == 0 or global_step == self.config.train_steps:
                structural_loss = self._validate_structural_loss(pipe, images, target_object)
                avg_train_loss = running_loss / max(1, self.config.validation_interval)
                running_loss = 0.0
                LOGGER.info(
                    "step=%d/%d train_loss=%.6f validation_structural_loss=%.6f best=%.6f elapsed=%.1fs",
                    global_step,
                    self.config.train_steps,
                    avg_train_loss,
                    structural_loss,
                    best_structural_loss,
                    time.time() - start,
                )
                if structural_loss + self.config.min_delta < best_structural_loss:
                    best_structural_loss = structural_loss
                    patience_counter = 0
                    self._save_lora(pipe, output_dir)
                    LOGGER.info("Saved improved LoRA checkpoint to %s", output_dir)
                else:
                    patience_counter += 1
                    if patience_counter >= self.config.patience:
                        LOGGER.info(
                            "Early stopping triggered at step %d after %d stagnant validations.",
                            global_step,
                            patience_counter,
                        )
                        break

        self._save_lora(pipe, output_dir)
        metadata = {
            "target_object": target_object,
            "instance_token": self.config.instance_token,
            "pretrained_model": self.config.pretrained_model,
            "best_structural_loss": best_structural_loss,
            "steps_completed": global_step,
            "training_config": asdict(self.config),
        }
        (output_dir / "conditioning_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
        LOGGER.info("Training complete. LoRA artifacts are available in %s", output_dir)
        return output_dir

    def _validate_structural_loss(self, pipe: Any, references: Sequence[Image.Image], target_object: str) -> float:
        pipe.unet.eval()
        if self.config.train_text_encoder_lora:
            pipe.text_encoder.eval()
        prompt = f"a centered studio product photo of {self.config.instance_token} {target_object}, neutral background, crisp outline"
        generator = torch.Generator(device=self.device).manual_seed(self.config.seed + 17)
        losses: List[float] = []
        try:
            with torch.no_grad():
                generated = pipe(
                    prompt=prompt,
                    negative_prompt=self.config.negative_prompt,
                    num_images_per_prompt=self.config.num_validation_images,
                    num_inference_steps=18,
                    guidance_scale=6.0,
                    height=self.config.resolution,
                    width=self.config.resolution,
                    generator=generator,
                ).images
            for image in generated:
                losses.append(compute_structural_loss(image, references, self.config.resolution))
        except torch.cuda.OutOfMemoryError as exc:
            clear_vram()
            raise RuntimeError("CUDA VRAM exhausted during structural validation.") from exc
        finally:
            pipe.unet.train()
            if self.config.train_text_encoder_lora:
                pipe.text_encoder.train()
        if not losses:
            return float("inf")
        return float(np.mean(losses))

    def _save_lora(self, pipe: Any, output_dir: Path) -> None:
        output_dir.mkdir(parents=True, exist_ok=True)
        if hasattr(pipe, "save_lora_weights") and get_peft_model_state_dict is not None:
            save_kwargs: Dict[str, Any] = {"save_directory": str(output_dir), "unet_lora_layers": get_peft_model_state_dict(pipe.unet)}
            if self.config.train_text_encoder_lora:
                save_kwargs["text_encoder_lora_layers"] = get_peft_model_state_dict(pipe.text_encoder)
            pipe.save_lora_weights(**save_kwargs)
        elif hasattr(pipe.unet, "save_pretrained"):
            pipe.unet.save_pretrained(str(output_dir / "unet_lora"))
        else:
            raise RuntimeError("The loaded pipeline cannot save LoRA weights with the installed Diffusers version.")


class ParametricSynthesizer:
    """Prompt-conditioned synthetic frame generator with systematic visual variation."""

    def __init__(self, config: SynthesisConfig, instance_token: str = "sksobj") -> None:
        self.config = config
        self.instance_token = instance_token
        self.device = resolve_device()
        self.dtype = torch.float32 if self.device == "cuda" else torch.float32
        self.pipe: Optional[Any] = None

    def _load_pipeline(self) -> Any:
        require_dependency("diffusers", _DIFFUSERS_IMPORT_ERROR)
        LOGGER.info("Loading synthesis pipeline: %s", self.config.pretrained_model)
        try:
            is_sdxl = "xl" in self.config.pretrained_model.lower() or "sdxl" in self.config.pretrained_model.lower()
            pipeline_cls = StableDiffusionXLPipeline if is_sdxl else StableDiffusionPipeline
            load_kwargs: Dict[str, Any] = {"torch_dtype": self.dtype}
            if not is_sdxl:
                load_kwargs.update({"safety_checker": None, "requires_safety_checker": False})
            pipe = pipeline_cls.from_pretrained(self.config.pretrained_model, **load_kwargs)
            pipe.to(self.device)
            if hasattr(pipe, "enable_attention_slicing"):
                pipe.enable_attention_slicing()
            if hasattr(pipe, "enable_vae_slicing"):
                pipe.enable_vae_slicing()
            if self.config.lora_dir:
                lora_path = Path(self.config.lora_dir).expanduser().resolve()
                if lora_path.exists():
                    pipe.load_lora_weights(str(lora_path))
                    if hasattr(pipe, "set_adapters"):
                        try:
                            pipe.set_adapters(["default_0"], adapter_weights=[self.config.lora_scale])
                        except Exception:
                            LOGGER.debug("Adapter weighting API unavailable or adapter name differs; using loaded LoRA default scale.")
                    LOGGER.info("Loaded LoRA weights from %s", lora_path)
                else:
                    raise FileNotFoundError(f"Configured LoRA directory does not exist: {lora_path}")
            self.pipe = pipe
            return pipe
        except torch.cuda.OutOfMemoryError as exc:
            clear_vram()
            raise RuntimeError("CUDA VRAM exhausted while loading the synthesis pipeline.") from exc
        except Exception as exc:
            clear_vram()
            raise RuntimeError(f"Failed to load synthesis pipeline: {exc}") from exc

    def generate(self) -> List[Tuple[Path, Image.Image, str]]:
        pipe = self.pipe or self._load_pipeline()
        output_root = Path(self.config.output_dir).expanduser().resolve()
        image_dir = output_root / "images"
        label_dir = output_root / "labels"
        image_dir.mkdir(parents=True, exist_ok=True)
        label_dir.mkdir(parents=True, exist_ok=True)
        generated_records: List[Tuple[Path, Image.Image, str]] = []
        LOGGER.info("Generating %d synthetic frames into %s", self.config.num_images, image_dir)
        for start_idx in range(0, self.config.num_images, self.config.batch_size):
            current_batch = min(self.config.batch_size, self.config.num_images - start_idx)
            prompts = [
                create_prompt_variation(self.config.prompt, self.config.target_object, self.instance_token, start_idx + i)
                for i in range(current_batch)
            ]
            generators = [torch.Generator(device=self.device).manual_seed(self.config.seed + start_idx + i) for i in range(current_batch)]
            try:
                with torch.no_grad():
                    result = pipe(
                        prompt=prompts,
                        negative_prompt=[self.config.negative_prompt] * current_batch,
                        width=self.config.width,
                        height=self.config.height,
                        num_inference_steps=self.config.inference_steps,
                        guidance_scale=self.config.guidance_scale,
                        generator=generators,
                    )
            except torch.cuda.OutOfMemoryError as exc:
                clear_vram()
                raise RuntimeError("CUDA VRAM exhausted during synthesis. Lower batch size, resolution, or steps.") from exc
            for local_idx, image in enumerate(result.images):
                image_id = start_idx + local_idx
                post_image = self._postprocess_variation(image, image_id)
                image_path = image_dir / f"synthetic_{image_id:06d}.jpg"
                post_image.save(image_path, quality=95)
                generated_records.append((image_path, post_image, prompts[local_idx]))
                LOGGER.debug("Generated %s with prompt: %s", image_path.name, prompts[local_idx])
        LOGGER.info("Generated %d images.", len(generated_records))
        return generated_records

    def _postprocess_variation(self, image: Image.Image, index: int) -> Image.Image:
        rng = random.Random(self.config.seed + index * 1543)
        image = image.convert("RGB")
        if rng.random() < 0.35:
            overlay = Image.new("RGB", image.size, (255, 255, 255))
            alpha = rng.uniform(0.015, 0.06)
            image = Image.blend(image, overlay, alpha)
        if rng.random() < 0.35:
            image = ImageEnhance.Brightness(image).enhance(rng.uniform(0.88, 1.12))
        if rng.random() < 0.35:
            image = ImageEnhance.Contrast(image).enhance(rng.uniform(0.9, 1.15))
        if rng.random() < 0.25:
            image = image.filter(ImageFilter.GaussianBlur(radius=rng.uniform(0.0, 0.45)))
        return image


class GroundingDINOLabeler:
    """Grounding DINO wrapper for zero-shot bounding-box extraction."""

    def __init__(self, model_id: str, box_threshold: float, text_threshold: float) -> None:
        require_dependency("transformers", _TRANSFORMERS_IMPORT_ERROR)
        require_dependency("torch", _TORCH_IMPORT_ERROR)
        self.model_id = model_id
        self.box_threshold = box_threshold
        self.text_threshold = text_threshold
        self.device = resolve_device()
        LOGGER.info("Loading Grounding DINO detector: %s", model_id)
        try:
            self.processor = AutoProcessor.from_pretrained(model_id)
            self.model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(self.device)
            self.model.eval()
        except torch.cuda.OutOfMemoryError as exc:
            clear_vram()
            raise RuntimeError("CUDA VRAM exhausted while loading Grounding DINO.") from exc
        except Exception as exc:
            clear_vram()
            raise RuntimeError(f"Failed to load Grounding DINO model '{model_id}': {exc}") from exc

    def detect(self, image: Image.Image, query: str, image_id: int) -> List[DetectionRecord]:
        text_labels = [[query]]
        width, height = image.size
        try:
            inputs = self.processor(images=image, text=text_labels, return_tensors="pt").to(self.device)
            with torch.no_grad():
                outputs = self.model(**inputs)
            results = self.processor.post_process_grounded_object_detection(
                outputs,
                inputs.input_ids,
                threshold=self.box_threshold,
                text_threshold=self.text_threshold,
                target_sizes=[(height, width)],
            )[0]
        except torch.cuda.OutOfMemoryError as exc:
            clear_vram()
            raise RuntimeError("CUDA VRAM exhausted during Grounding DINO inference.") from exc
        except Exception as exc:
            raise RuntimeError(f"Grounding DINO inference failed for image_id={image_id}: {exc}") from exc

        records: List[DetectionRecord] = []
        boxes = results.get("boxes", [])
        scores = results.get("scores", [])
        labels = results.get("labels", [])
        for box, score, label in zip(boxes, scores, labels):
            box_tuple = tuple(float(x) for x in box.detach().cpu().tolist())
            score_float = float(score.detach().cpu().item()) if hasattr(score, "detach") else float(score)
            label_text = str(label)
            rec = DetectionRecord(image_id=image_id, label=label_text, score=score_float, box_xyxy=box_tuple, width=width, height=height).clipped()
            xmin, ymin, xmax, ymax = rec.box_xyxy
            if xmax - xmin >= 2 and ymax - ymin >= 2:
                records.append(rec)
        LOGGER.debug("Detector returned %d boxes for image_id=%d", len(records), image_id)
        return records


def convert_detection(record: DetectionRecord, fmt: LabelFormat, class_id: int = 0, category_id: int = 1) -> Union[List[float], Dict[str, Any]]:
    rec = record.clipped()
    xmin, ymin, xmax, ymax = rec.box_xyxy
    box_w = max(0.0, xmax - xmin)
    box_h = max(0.0, ymax - ymin)
    if fmt == "yolo":
        x_center = (xmin + box_w / 2.0) / rec.width
        y_center = (ymin + box_h / 2.0) / rec.height
        return [
            int(class_id),
            round(float(x_center), 6),
            round(float(y_center), 6),
            round(float(box_w / rec.width), 6),
            round(float(box_h / rec.height), 6),
        ]
    if fmt == "coco":
        return {
            "image_id": int(rec.image_id),
            "category_id": int(category_id),
            "bbox": [round(float(xmin), 2), round(float(ymin), 2), round(float(box_w), 2), round(float(box_h), 2)],
            "score": round(float(rec.score), 6),
            "label": rec.label,
        }
    raise ValueError(f"Unsupported label format: {fmt}")


def write_label_file(label_dir: Path, image_path: Path, detections: Sequence[DetectionRecord], config: SynthesisConfig) -> Path:
    label_dir.mkdir(parents=True, exist_ok=True)
    if config.label_format == "yolo":
        label_path = label_dir / f"{image_path.stem}.txt"
        lines = []
        for rec in detections:
            converted = convert_detection(rec, "yolo", class_id=config.class_id)
            lines.append(" ".join(str(x) for x in converted))
        label_path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")
        return label_path
    label_path = label_dir / f"{image_path.stem}.json"
    records = [convert_detection(rec, "coco", category_id=config.category_id) for rec in detections]
    label_path.write_text(json.dumps(records, indent=2), encoding="utf-8")
    return label_path


def draw_debug_previews(
    output_dir: Union[str, Path],
    image_records: Sequence[Tuple[Path, List[DetectionRecord]]],
    label_format: LabelFormat,
    count: int = 5,
    seed: int = 1337,
) -> List[Path]:
    require_dependency("opencv-python", _CV2_IMPORT_ERROR)
    output_root = Path(output_dir).expanduser().resolve()
    if not image_records:
        LOGGER.warning("No image records available for debug preview generation.")
        return []
    rng = random.Random(seed)
    sample = list(image_records)
    rng.shuffle(sample)
    selected = sample[: min(count, len(sample))]
    preview_paths: List[Path] = []
    for idx, (image_path, detections) in enumerate(selected):
        img = cv2.imread(str(image_path))
        if img is None:
            LOGGER.warning("OpenCV could not read image for debug preview: %s", image_path)
            continue
        for rec in detections:
            xmin, ymin, xmax, ymax = [int(round(v)) for v in rec.clipped().box_xyxy]
            cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (40, 220, 40), 2)
            label = f"{rec.label} {rec.score:.2f}"
            cv2.putText(img, label, (xmin, max(15, ymin - 6)), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (40, 220, 40), 1, cv2.LINE_AA)
        preview_path = output_root / f"debug_preview_{idx:02d}.jpg"
        cv2.imwrite(str(preview_path), img)
        preview_paths.append(preview_path)
        LOGGER.info("Saved %s debug preview: %s", label_format.upper(), preview_path)
    return preview_paths


def compile_and_label_outputs(generated: Sequence[Tuple[Path, Image.Image, str]], config: SynthesisConfig) -> Dict[str, Any]:
    output_root = Path(config.output_dir).expanduser().resolve()
    image_dir = output_root / "images"
    label_dir = output_root / "labels"
    image_dir.mkdir(parents=True, exist_ok=True)
    label_dir.mkdir(parents=True, exist_ok=True)
    labeler = GroundingDINOLabeler(config.detector_model, config.detection_threshold, config.text_threshold)
    image_records: List[Tuple[Path, List[DetectionRecord]]] = []
    total_boxes = 0
    for image_id, (image_path, image, prompt) in enumerate(generated):
        detections = labeler.detect(image, query=config.target_object, image_id=image_id)
        write_label_file(label_dir, image_path, detections, config)
        image_records.append((image_path, detections))
        total_boxes += len(detections)
        LOGGER.info("Labeled image_id=%d file=%s boxes=%d", image_id, image_path.name, len(detections))
    previews = draw_debug_previews(output_root, image_records, config.label_format, config.debug_preview_count, config.seed)
    manifest = {
        "output_dir": str(output_root),
        "image_dir": str(image_dir),
        "label_dir": str(label_dir),
        "label_format": config.label_format,
        "target_object": config.target_object,
        "num_images": len(generated),
        "total_boxes": total_boxes,
        "debug_previews": [str(path) for path in previews],
        "config": asdict(config),
    }
    (output_root / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    LOGGER.info("Compiled labeled dataset in %s with %d boxes.", output_root, total_boxes)
    return manifest


def run_pipeline(source_images: Sequence[ImageLike], train_cfg: TrainingConfig, synth_cfg: SynthesisConfig) -> Dict[str, Any]:
    LOGGER.info("Starting end-to-end synthetic data pipeline.")
    trainer = FewShotLoRATrainer(train_cfg)
    lora_dir = trainer.train(source_images, synth_cfg.target_object)
    synth_cfg.lora_dir = str(lora_dir)
    synthesizer = ParametricSynthesizer(synth_cfg, instance_token=train_cfg.instance_token)
    generated = synthesizer.generate()
    manifest = compile_and_label_outputs(generated, synth_cfg)
    LOGGER.info("Pipeline complete: %s", manifest["output_dir"])
    return manifest


if modal is not None:  # pragma: no cover - only active in Modal runtime
    modal_image = (
        modal.Image.debian_slim(python_version="3.11")
        .pip_install(
            "torch",
            "diffusers",
            "transformers",
            "accelerate",
            "opencv-python-headless",
            "pillow",
            "peft",
            "safetensors",
        )
    )
    app = modal.App("standalone-cv-synthetic-data-engine", image=modal_image)

    @app.function(gpu="A10G", timeout=60 * 60 * 3)
    def modal_run_pipeline(source_dir: str, training_config: Dict[str, Any], synthesis_config: Dict[str, Any]) -> Dict[str, Any]:
        configure_logging("INFO")
        paths = load_image_paths(source_dir)
        train_cfg = TrainingConfig(**training_config)
        synth_cfg = SynthesisConfig(**synthesis_config)
        return run_pipeline(paths, train_cfg, synth_cfg)
else:
    app = None


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Few-shot synthetic CV data generation pipeline.")
    parser.add_argument("--source-dir", required=True, help="Directory containing target object source images.")
    parser.add_argument("--target-object", required=True, help="Exact semantic text string for the target object detector query.")
    parser.add_argument("--prompt", required=True, help="Background/environment prompt, e.g. 'industrial conveyor belt with reflections'.")
    parser.add_argument("--output-dir", default="output_batch", help="Output dataset directory.")
    parser.add_argument("--pretrained-model", default="runwayml/stable-diffusion-v1-5", help="Open diffusion model identifier or local path.")
    parser.add_argument("--detector-model", default="IDEA-Research/grounding-dino-tiny", help="Grounding DINO model identifier.")
    parser.add_argument("--label-format", choices=["yolo", "coco"], default="yolo", help="Annotation output format.")
    parser.add_argument("--num-images", type=int, default=24, help="Number of synthetic images to generate.")
    parser.add_argument("--train-steps", type=int, default=180, help="Maximum LoRA training steps before early stopping.")
    parser.add_argument("--resolution", type=int, default=512, help="Training image resolution.")
    parser.add_argument("--width", type=int, default=512, help="Generated image width.")
    parser.add_argument("--height", type=int, default=512, help="Generated image height.")
    parser.add_argument("--batch-size", type=int, default=1, help="Synthesis batch size.")
    parser.add_argument("--train-batch-size", type=int, default=1, help="LoRA training batch size.")
    parser.add_argument("--rank", type=int, default=8, help="LoRA rank.")
    parser.add_argument("--learning-rate", type=float, default=1e-4, help="LoRA learning rate.")
    parser.add_argument("--seed", type=int, default=1337, help="Random seed.")
    parser.add_argument("--log-level", default="INFO", help="Logging verbosity.")
    return parser.parse_args(argv)


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)
    configure_logging(args.log_level)
    try:
        source_paths = load_image_paths(args.source_dir)
        lora_dir = str(Path(args.output_dir).expanduser().resolve() / "conditioned_lora")
        train_cfg = TrainingConfig(
            pretrained_model=args.pretrained_model,
            output_dir=lora_dir,
            resolution=args.resolution,
            train_steps=args.train_steps,
            batch_size=args.train_batch_size,
            rank=args.rank,
            learning_rate=args.learning_rate,
            seed=args.seed,
        )
        synth_cfg = SynthesisConfig(
            prompt=args.prompt,
            target_object=args.target_object,
            output_dir=args.output_dir,
            lora_dir=lora_dir,
            pretrained_model=args.pretrained_model,
            num_images=args.num_images,
            batch_size=args.batch_size,
            width=args.width,
            height=args.height,
            seed=args.seed,
            label_format=args.label_format,
            detector_model=args.detector_model,
        )
        manifest = run_pipeline(source_paths, train_cfg, synth_cfg)
        LOGGER.info("Final manifest: %s", json.dumps(manifest, indent=2))
        return 0
    except KeyboardInterrupt:
        LOGGER.warning("Pipeline interrupted by user.")
        return 130
    except Exception as exc:
        LOGGER.exception("Pipeline failed: %s", exc)
        return 1
    finally:
        clear_vram()


if __name__ == "__main__":
    raise SystemExit(main())