diffusers-workflow 0.4.0a3__py3-none-any.whl

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