diffusers-workflow 0.4.0__py3-none-any.whl

dw/tasks/restore_faces.py

@@ -0,0 +1,175 @@
+"""
+Face restoration via spandrel + facexlib.
+
+Uses facexlib for face detection/alignment/pasting and spandrel for
+neural network inference on cropped faces. Supports GFPGAN, RestoreFormer,
+and CodeFormer (via spandrel-extra-arches) model weights.
+"""
+
+import logging
+import numpy as np
+import torch
+from PIL import Image
+
+from .model_cache import cached_model
+
+logger = logging.getLogger("dw")
+
+
+def restore_faces(image, model_name, device="cpu", **kwargs):
+    """Restore faces in an image using a spandrel-compatible face restoration model.
+
+    Args:
+        image: PIL Image containing faces to restore
+        model_name: HuggingFace repo ID or local path to model weights
+        device: Target device ("cuda", "mps", "cpu")
+        **kwargs:
+            filename: Weight file name within a HF repo (default: auto-detect)
+            upscale_factor: Background upscale factor (default: 1, no upscaling)
+            face_size: Cropped face size in pixels (default: 512)
+            use_parse: Use face parsing for better blending (default: True)
+            only_center_face: Only restore the largest/center face (default: False)
+            detection_resize: Resize shorter side for detection speed (default: 640)
+            eye_dist_threshold: Skip faces with eye distance below this (default: 5)
+            upsample_img: Pre-upscaled background PIL Image (default: None)
+
+    Returns:
+        PIL Image with restored faces
+    """
+    try:
+        from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+    except ImportError:
+        raise ImportError(
+            "facexlib is required for face restoration. Install with: pip install facexlib"
+        )
+
+    from .upscale import _resolve_model_path
+
+    filename = kwargs.get("filename", None)
+    upscale_factor = kwargs.get("upscale_factor", 1)
+    face_size = kwargs.get("face_size", 512)
+    use_parse = kwargs.get("use_parse", True)
+    only_center_face = kwargs.get("only_center_face", False)
+    detection_resize = kwargs.get("detection_resize", 640)
+    eye_dist_threshold = kwargs.get("eye_dist_threshold", 5)
+    upsample_img = kwargs.get("upsample_img", None)
+
+    # Load the face restoration model via spandrel
+    def load_descriptor():
+        model_path = _resolve_model_path(model_name, filename)
+        result = _load_face_model(model_path, device)
+        if device != "cpu" and result.supports_half:
+            result.model.half()
+        return result
+
+    descriptor = cached_model(
+        ("restore_faces", model_name, filename, str(device)), load_descriptor
+    )
+
+    # Set up facexlib helper
+    face_helper = FaceRestoreHelper(
+        upscale_factor=upscale_factor,
+        face_size=face_size,
+        crop_ratio=(1, 1),
+        det_model="retinaface_resnet50",
+        use_parse=use_parse,
+        device=torch.device(device),
+    )
+
+    # Convert PIL to BGR numpy (facexlib format)
+    input_bgr = np.array(image.convert("RGB"))[:, :, ::-1].copy()
+    face_helper.read_image(input_bgr)
+
+    # Detect faces
+    num_faces = face_helper.get_face_landmarks_5(
+        only_center_face=only_center_face,
+        resize=detection_resize,
+        eye_dist_threshold=eye_dist_threshold,
+    )
+    logger.info(f"Detected {num_faces} face(s)")
+
+    if num_faces == 0:
+        logger.warning("No faces detected, returning original image")
+        return image
+
+    # Align and warp faces to face_size x face_size
+    face_helper.align_warp_face()
+
+    # Half precision was applied at load time (see load_descriptor above) if supported
+    use_half = device != "cpu" and descriptor.supports_half
+    model_dtype = torch.float16 if use_half else torch.float32
+
+    # Restore each face
+    for i, cropped_face in enumerate(face_helper.cropped_faces):
+        logger.debug(f"Restoring face {i + 1}/{num_faces}")
+
+        # BGR uint8 numpy -> float32 tensor [1, 3, H, W]
+        face_tensor = (
+            torch.from_numpy(cropped_face.astype(np.float32) / 255.0)
+            .permute(2, 0, 1)
+            .unsqueeze(0)
+        )
+        face_tensor = face_tensor.to(device=device, dtype=model_dtype)
+
+        with torch.inference_mode():
+            restored_tensor = descriptor(face_tensor)
+
+        # Tensor -> BGR uint8 numpy. Rounds rather than truncates when
+        # quantizing (matches diffusers' VaeImageProcessor.numpy_to_pil
+        # behavior) so exact 8-bit values don't drift down a level; this
+        # stays local rather than using tensor_image's shared helpers
+        # because the data here is BGR numpy (facexlib's format), not PIL/RGB.
+        restored = restored_tensor.squeeze(0).permute(1, 2, 0)
+        restored = restored.mul(255).round().clamp(0, 255).byte().cpu().numpy()
+
+        # Resize to expected face_size if model output differs
+        if restored.shape[:2] != (face_size, face_size):
+            import cv2
+
+            restored = cv2.resize(
+                restored, (face_size, face_size), interpolation=cv2.INTER_LANCZOS4
+            )
+
+        face_helper.add_restored_face(restored)
+
+    # Prepare inverse affine transforms
+    face_helper.get_inverse_affine()
+
+    # Paste faces back onto the image
+    upsample_bgr = None
+    if upsample_img is not None:
+        upsample_bgr = np.array(upsample_img.convert("RGB"))[:, :, ::-1].copy()
+
+    result_bgr = face_helper.paste_faces_to_input_image(upsample_img=upsample_bgr)
+
+    # BGR numpy -> PIL RGB
+    result = Image.fromarray(result_bgr[:, :, ::-1])
+    logger.info(
+        f"Face restoration complete ({num_faces} face(s), {result.width}x{result.height})"
+    )
+
+    face_helper.clean_all()
+    return result
+
+
+def _load_face_model(model_path, device):
+    """Load a face restoration model via spandrel."""
+    try:
+        from spandrel import ModelLoader, ImageModelDescriptor
+    except ImportError:
+        raise ImportError(
+            "spandrel is required for face restoration. Install with: pip install spandrel"
+        )
+
+    logger.info(f"Loading face restoration model from {model_path}")
+    loader = ModelLoader(device=torch.device(device))
+    descriptor = loader.load_from_file(model_path)
+
+    if not isinstance(descriptor, ImageModelDescriptor):
+        raise ValueError(
+            f"Model is not an image model (got {type(descriptor).__name__}). "
+            f"Expected a face restoration model (GFPGAN, CodeFormer, RestoreFormer)."
+        )
+
+    logger.info(f"Loaded {descriptor.architecture.name}")
+    return descriptor

dw/tasks/rife_model.py

@@ -0,0 +1,192 @@
+"""
+RIFE IFNet v4.13 — Real-Time Intermediate Flow Estimation.
+
+Vendored from https://github.com/hzwer/Practical-RIFE
+Original architecture by Zhewei Huang et al.
+Matches the v4.13.2 checkpoint layout (feature-encode Head + four IFBlocks).
+
+MIT License — Copyright (c) Megvii Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def warp(tenInput, tenFlow, tenFlow_div, backwarp_tenGrid):
+    """Backward warp tenInput according to optical flow tenFlow."""
+    tenFlow = torch.cat(
+        [
+            tenFlow[:, 0:1] / tenFlow_div[0],
+            tenFlow[:, 1:2] / tenFlow_div[1],
+        ],
+        1,
+    )
+    g = (backwarp_tenGrid + tenFlow).permute(0, 2, 3, 1)
+    return F.grid_sample(
+        input=tenInput,
+        grid=g,
+        mode="bilinear",
+        padding_mode="border",
+        align_corners=True,
+    )
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=True,
+        ),
+        nn.LeakyReLU(0.2, True),
+    )
+
+
+class ResConv(nn.Module):
+    def __init__(self, c, dilation=1):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            c, c, 3, 1, dilation, dilation=dilation, groups=1, bias=True
+        )
+        self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True)
+        self.relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x):
+        return self.relu(self.conv(x) * self.beta + x)
+
+
+class Head(nn.Module):
+    """Feature encoder — 8-channel per-image features consumed by the IFBlocks."""
+
+    def __init__(self):
+        super().__init__()
+        self.cnn0 = nn.Conv2d(3, 32, 3, 2, 1)
+        self.cnn1 = nn.Conv2d(32, 32, 3, 1, 1)
+        self.cnn2 = nn.Conv2d(32, 32, 3, 1, 1)
+        self.cnn3 = nn.ConvTranspose2d(32, 8, 4, 2, 1)
+        self.relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x):
+        x = self.relu(self.cnn0(x))
+        x = self.relu(self.cnn1(x))
+        x = self.relu(self.cnn2(x))
+        return self.cnn3(x)
+
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, c=64):
+        super().__init__()
+        self.conv0 = nn.Sequential(
+            conv(in_planes, c // 2, 3, 2, 1),
+            conv(c // 2, c, 3, 2, 1),
+        )
+        self.convblock = nn.Sequential(
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+        )
+        self.lastconv = nn.Sequential(
+            nn.ConvTranspose2d(c, 4 * 6, 4, 2, 1),
+            nn.PixelShuffle(2),
+        )
+
+    def forward(self, x, flow=None, scale=1):
+        x = F.interpolate(
+            x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
+        )
+        if flow is not None:
+            flow = F.interpolate(
+                flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
+            ) * (1.0 / scale)
+            x = torch.cat((x, flow), 1)
+        feat = self.conv0(x)
+        feat = self.convblock(feat)
+        tmp = self.lastconv(feat)
+        tmp = F.interpolate(
+            tmp, scale_factor=scale, mode="bilinear", align_corners=False
+        )
+        flow = tmp[:, :4] * scale
+        mask = tmp[:, 4:5]
+        return flow, mask
+
+
+class IFNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.block0 = IFBlock(7 + 16, c=192)
+        self.block1 = IFBlock(8 + 4 + 16, c=128)
+        self.block2 = IFBlock(8 + 4 + 16, c=96)
+        self.block3 = IFBlock(8 + 4 + 16, c=64)
+        self.encode = Head()
+
+    def forward(self, img0, img1, timestep, scale_list, tenFlow_div, backwarp_tenGrid):
+        f0 = self.encode(img0[:, :3])
+        f1 = self.encode(img1[:, :3])
+        flow_list = []
+        merged = []
+        mask_list = []
+        warped_img0 = img0
+        warped_img1 = img1
+
+        flow = None
+        for i, (block, scale) in enumerate(
+            zip(
+                [self.block0, self.block1, self.block2, self.block3],
+                scale_list,
+            )
+        ):
+            if flow is None:
+                flow, mask = block(
+                    torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1),
+                    None,
+                    scale=scale,
+                )
+            else:
+                wf0 = warp(f0, flow[:, :2], tenFlow_div, backwarp_tenGrid)
+                wf1 = warp(f1, flow[:, 2:4], tenFlow_div, backwarp_tenGrid)
+                fd, m = block(
+                    torch.cat((warped_img0, warped_img1, wf0, wf1, timestep, mask), 1),
+                    flow,
+                    scale=scale,
+                )
+                flow = flow + fd
+                mask = mask + m
+
+            mask_list.append(mask)
+            flow_list.append(flow)
+            warped_img0 = warp(img0[:, :3], flow[:, :2], tenFlow_div, backwarp_tenGrid)
+            warped_img1 = warp(img1[:, :3], flow[:, 2:4], tenFlow_div, backwarp_tenGrid)
+            merged.append(
+                warped_img0 * torch.sigmoid(mask)
+                + warped_img1 * (1 - torch.sigmoid(mask))
+            )
+
+        return flow_list, mask_list, merged

dw/tasks/segment.py

@@ -0,0 +1,121 @@
+"""
+Image segmentation via GroundingDINO + SAM2.
+
+Takes an image and a text prompt, detects objects matching the prompt,
+and returns a binary mask image (white = detected object).
+"""
+
+import logging
+import torch
+import numpy as np
+from PIL import Image, ImageOps
+from transformers import (
+    AutoProcessor,
+    AutoModelForZeroShotObjectDetection,
+    Sam2Model,
+    Sam2Processor,
+)
+from .model_cache import cached_model
+
+logger = logging.getLogger("dw")
+
+_DEFAULT_DINO_MODEL = "IDEA-Research/grounding-dino-base"
+_DEFAULT_SAM_MODEL = "facebook/sam2-hiera-large"
+
+
+def segment_image(image, prompt, device="cpu", **kwargs):
+    """Segment objects matching a text prompt, returning a binary mask.
+
+    Args:
+        image: PIL Image to segment
+        prompt: Text description of object(s) to detect (e.g., "dog")
+        device: Target device ("cuda", "mps", "cpu")
+        **kwargs:
+            model_name: GroundingDINO model ID
+            sam_model_name: SAM2 model ID
+            threshold: Detection confidence threshold (default: 0.3)
+            invert: Invert the output mask (default: False)
+
+    Returns:
+        PIL Image in mode "L" — white (255) for detected objects, black (0) for background.
+    """
+    model_name = kwargs.get("model_name", _DEFAULT_DINO_MODEL)
+    sam_model_name = kwargs.get("sam_model_name", _DEFAULT_SAM_MODEL)
+    threshold = kwargs.get("threshold", 0.3)
+    invert = kwargs.get("invert", False)
+
+    width, height = image.size
+
+    def load_dino():
+        logger.info(f"Loading GroundingDINO from {model_name}")
+        processor = AutoProcessor.from_pretrained(model_name)
+        model = AutoModelForZeroShotObjectDetection.from_pretrained(model_name).to(
+            device
+        )
+        return processor, model
+
+    dino_processor, dino_model = cached_model(
+        ("segment_dino", model_name, str(device)), load_dino
+    )
+
+    inputs = dino_processor(images=image, text=prompt, return_tensors="pt").to(device)
+
+    with torch.inference_mode():
+        outputs = dino_model(**inputs)
+
+    results = dino_processor.post_process_grounded_object_detection(
+        outputs,
+        inputs["input_ids"],
+        threshold=threshold,
+        target_sizes=[(height, width)],
+    )
+
+    boxes = results[0]["boxes"]
+    scores = results[0]["scores"]
+    labels = results[0]["labels"]
+
+    logger.info(f"Detected {len(boxes)} objects: {labels} (scores: {scores.tolist()})")
+
+    if len(boxes) == 0:
+        mask_image = Image.new("L", (width, height), 0)
+        if invert:
+            mask_image = ImageOps.invert(mask_image)
+        return mask_image
+
+    def load_sam():
+        logger.info(f"Loading SAM2 from {sam_model_name}")
+        processor = Sam2Processor.from_pretrained(sam_model_name)
+        model = Sam2Model.from_pretrained(sam_model_name).to(device)
+        return processor, model
+
+    sam_processor, sam_model = cached_model(
+        ("segment_sam", sam_model_name, str(device)), load_sam
+    )
+
+    input_boxes = [boxes.cpu().tolist()]
+
+    sam_inputs = sam_processor(
+        images=image,
+        input_boxes=input_boxes,
+        return_tensors="pt",
+    ).to(device)
+
+    with torch.inference_mode():
+        sam_outputs = sam_model(**sam_inputs)
+
+    masks = sam_processor.post_process_masks(
+        sam_outputs.pred_masks,
+        sam_inputs["original_sizes"],
+        sam_inputs["reshaped_input_sizes"],
+    )
+
+    combined = masks[0][:, 0].sum(dim=0).clamp(0, 1)
+    mask_array = (combined.cpu().numpy() * 255).astype(np.uint8)
+
+    mask_image = Image.fromarray(mask_array, mode="L")
+
+    if invert:
+        mask_image = ImageOps.invert(mask_image)
+
+    logger.info(f"Generated segmentation mask {width}x{height}")
+    return mask_image

dw/tasks/select.py

@@ -0,0 +1,111 @@
+"""Reduce a list of candidates to one by a deterministic rule.
+
+The README's headline shape is N candidates -> choose one -> spend an
+expensive stage on the winner. Today that choice lives at the agent
+boundary (keep_output + asset:); this task is the version of it that needs
+no judgment - argmax, a threshold, or a fixed index - so the recipe stays
+inside the workflow and replays. See docs/proposals/score-and-select.md.
+"""
+
+import logging
+
+logger = logging.getLogger("dw")
+
+
+class Selected:
+    """The winning candidate, plus the metadata that makes the choice
+    replayable (position, score). Compares equal to its own value so a
+    caller that only wants the winner can treat it as one."""
+
+    def __init__(self, value, position, score):
+        self.value = value
+        self.position = position
+        self.score = score
+
+    def __eq__(self, other):
+        if isinstance(other, Selected):
+            return (
+                self.value == other.value
+                and self.position == other.position
+                and self.score == other.score
+            )
+        return self.value == other
+
+    def __hash__(self):
+        return hash(self.value)
+
+    def __repr__(self):
+        return f"Selected(value={self.value!r}, position={self.position}, score={self.score!r})"
+
+
+_THRESHOLD_RULES = {
+    "first_above": lambda score, threshold: score >= threshold,
+    "first_below": lambda score, threshold: score <= threshold,
+}
+
+
+def _parse_score(index, score):
+    if isinstance(score, bool) or not isinstance(score, (str, int, float)):
+        raise ValueError(f"select: score {index} is not a number: {score!r}")
+    try:
+        return float(score)
+    except (TypeError, ValueError):
+        raise ValueError(f"select: score {index} is not a number: {score!r}")
+
+
+def select(candidates, scores, rule, threshold=None, index=None):
+    """Task command: pick one candidate out of many by a fixed rule.
+
+    Args:
+        candidates: The values to choose among, in order.
+        scores: One number (or numeric string) per candidate, same length.
+        rule: "argmax" | "argmin" | "first_above" | "first_below" | "index"
+        threshold: Required by first_above/first_below.
+        index: Required by rule "index" - a position into candidates.
+
+    Returns:
+        A Selected wrapping the winning value, its position and its score.
+
+    Raises:
+        ValueError: length mismatch, a non-numeric score, an unknown rule,
+            a missing threshold/index, an out-of-range index, or (for
+            first_above/first_below) no candidate meeting the threshold.
+    """
+    if len(candidates) != len(scores):
+        raise ValueError(
+            f"select got {len(candidates)} candidates and {len(scores)} scores "
+            "- candidates and scores must be the same length"
+        )
+
+    parsed_scores = [_parse_score(i, s) for i, s in enumerate(scores)]
+
+    if rule in ("argmax", "argmin"):
+        positions = range(len(parsed_scores))
+        picker = max if rule == "argmax" else min
+        position = picker(positions, key=lambda i: parsed_scores[i])
+    elif rule in _THRESHOLD_RULES:
+        if threshold is None:
+            raise ValueError(f"select rule '{rule}' requires a threshold")
+        passes = _THRESHOLD_RULES[rule]
+        position = next(
+            (i for i, s in enumerate(parsed_scores) if passes(s, threshold)), None
+        )
+        if position is None:
+            raise ValueError(
+                f"select: no candidate passes rule '{rule}' at threshold {threshold}"
+            )
+    elif rule == "index":
+        if index is None:
+            raise ValueError("select rule 'index' requires an index")
+        if index < 0 or index >= len(candidates):
+            raise ValueError(
+                f"select: index {index} is out of range for {len(candidates)} candidates"
+            )
+        position = index
+    else:
+        raise ValueError(f"select: unknown rule: {rule!r}")
+
+    logger.debug(f"select: rule={rule} chose position {position}")
+    return Selected(
+        value=candidates[position], position=position, score=parsed_scores[position]
+    )