diffusers-workflow 0.4.0a3__py3-none-any.whl

dw/tasks/interpolate_frames.py

@@ -0,0 +1,227 @@
+"""
+Video frame interpolation via RIFE (Real-Time Intermediate Flow Estimation).
+
+Takes a list of video frames and generates intermediate frames to increase
+frame rate. Supports 2x, 4x, and 8x multipliers.
+
+Model weights are downloaded from HuggingFace Hub on first use.
+"""
+
+import logging
+import torch
+
+from .tensor_image import pil_to_float_tensor as _pil_to_tensor, float_tensor_to_pil
+
+logger = logging.getLogger("dw")
+
+_VALID_MULTIPLIERS = {2, 4, 8}
+
+
+def interpolate_frames(video, device="cpu", **kwargs):
+    """Interpolate between video frames using RIFE to increase frame rate.
+
+    Args:
+        video: List of PIL Images (video frames)
+        device: Target device ("cuda", "mps", "cpu")
+        **kwargs:
+            multiplier: Frame count multiplier — 2, 4, or 8 (default: 2)
+            model_name: HuggingFace repo with RIFE v4.13 weights (default: auto)
+            filename: Weights filename within the repo (default: auto)
+
+    Returns:
+        List of PIL Images with interpolated frames inserted.
+    """
+    multiplier = int(kwargs.get("multiplier", 2))
+    model_name = kwargs.get("model_name", None)
+    filename = kwargs.get("filename", None)
+
+    if multiplier not in _VALID_MULTIPLIERS:
+        raise ValueError(
+            f"multiplier must be one of {sorted(_VALID_MULTIPLIERS)}, got {multiplier}"
+        )
+
+    if len(video) < 2:
+        raise ValueError(f"Need at least 2 frames to interpolate, got {len(video)}")
+
+    logger.info(
+        f"Interpolating {len(video)} frames with {multiplier}x multiplier on {device}"
+    )
+
+    model = _load_rife_model(device, model_name, filename)
+
+    passes = {2: 1, 4: 2, 8: 3}[multiplier]
+    frames = list(video)
+
+    for pass_num in range(passes):
+        logger.debug(
+            f"Interpolation pass {pass_num + 1}/{passes}: {len(frames)} frames"
+        )
+        frames = _interpolate_2x(frames, model)
+
+    logger.info(f"Interpolation complete: {len(video)} -> {len(frames)} frames")
+    return frames
+
+
+def _interpolate_2x(frames, model):
+    """Single pass of 2x interpolation — insert one frame between each pair."""
+    result = [frames[0]]
+    for i in range(len(frames) - 1):
+        mid_frame = model(frames[i], frames[i + 1])
+        result.append(mid_frame)
+        result.append(frames[i + 1])
+    return result
+
+
+_DEFAULT_RIFE_REPO = "imaginairy/rife-interpolation"
+_DEFAULT_RIFE_FILENAME = "rife-flownet-4.13.2.safetensors"
+
+
+def _pad_tensor(t, ph, pw):
+    """Zero-pad a (1, 3, H, W) tensor on the bottom/right to (ph, pw)."""
+    h, w = t.shape[2], t.shape[3]
+    padding = (0, pw - w, 0, ph - h)
+    return torch.nn.functional.pad(t, padding)
+
+
+def _padded_size(h, w, multiple=32):
+    """Round (h, w) up to the next multiple (RIFE requires dims divisible by 32)."""
+    ph = ((h - 1) // multiple + 1) * multiple
+    pw = ((w - 1) // multiple + 1) * multiple
+    return ph, pw
+
+
+def _make_flow_context(ph, pw, device):
+    """Build the warp grid and flow divisors for one padded resolution."""
+    tenFlow_div = torch.tensor([(pw - 1.0) / 2.0, (ph - 1.0) / 2.0], device=device)
+    backwarp_tenGrid = torch.cat(
+        [
+            torch.linspace(-1.0, 1.0, pw, device=device)
+            .view(1, 1, 1, pw)
+            .expand(-1, -1, ph, -1),
+            torch.linspace(-1.0, 1.0, ph, device=device)
+            .view(1, 1, ph, 1)
+            .expand(-1, -1, -1, pw),
+        ],
+        1,
+    )
+    timestep = torch.full((1, 1, ph, pw), 0.5, dtype=torch.float32, device=device)
+    return tenFlow_div, backwarp_tenGrid, timestep
+
+
+def _load_rife_model(device, model_name=None, filename=None):
+    """Load RIFE model and return a callable that interpolates two frames.
+
+    Args:
+        device: Target device string ("cuda", "mps", "cpu").
+        model_name: Optional HuggingFace repo ID containing RIFE v4.13 weights.
+            Defaults to imaginairy/rife-interpolation.
+        filename: Optional weights filename within the repo. Defaults to
+            rife-flownet-4.13.2.safetensors. Both .safetensors and torch
+            checkpoint formats (.pkl/.pth) are supported.
+
+    Returns:
+        Callable that takes (frame1: PIL.Image, frame2: PIL.Image) -> PIL.Image
+    """
+    from .rife_model import IFNet
+    from huggingface_hub import hf_hub_download
+    from .model_cache import cached_model
+
+    repo_id = model_name if model_name is not None else _DEFAULT_RIFE_REPO
+    weights_file = filename if filename is not None else _DEFAULT_RIFE_FILENAME
+
+    def load_net():
+        model_path = hf_hub_download(repo_id=repo_id, filename=weights_file)
+
+        logger.info(f"Loading RIFE IFNet v4.13 to {device}")
+
+        if weights_file.endswith(".safetensors"):
+            from safetensors.torch import load_file
+
+            state_dict = load_file(model_path)
+        else:
+            state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
+
+        # Strip "module." prefix that comes from DataParallel-saved checkpoints
+        cleaned = {}
+        for k, v in state_dict.items():
+            cleaned[k.removeprefix("module.")] = v
+
+        net = IFNet()
+        net.load_state_dict(cleaned)
+        net.eval()
+        net.to(device)
+        return net
+
+    net = cached_model(
+        ("interpolate_frames", repo_id, weights_file, str(device)), load_net
+    )
+
+    return _build_inference(net, device)
+
+
+def _build_inference(net, device):
+    """Build the (img1, img2) -> mid_frame callable for a loaded RIFE net.
+
+    All frames of a video share one padded resolution, so the flow-warp
+    context (tenFlow_div, backwarp grid, timestep) is computed once per
+    (padded_h, padded_w) and reused for every pair instead of being rebuilt
+    on every call.
+
+    `_interpolate_2x` walks frames pairwise: (f0, f1), (f1, f2), (f2, f3), ...
+    — the second frame of one pair is the same PIL object as the first frame
+    of the next. This closure carries the padded tensor it produced for a
+    pair's second frame forward, so that when that exact frame object shows
+    up again as a pair's first frame, its tensor is reused instead of being
+    re-derived from the PIL image (re-decoded, re-normalized, re-padded). A
+    cache miss (non-matching object, e.g. across separate interpolation
+    passes) simply falls back to a fresh conversion, so this is a pure
+    optimization with no behavioral effect — the reused tensor is bit-for-bit
+    what a fresh conversion of the same frame would produce.
+    """
+    scale_list = [8, 4, 2, 1]
+    flow_context_cache = {}
+    carry = {"img": None}
+
+    def get_flow_context(ph, pw):
+        key = (ph, pw)
+        ctx = flow_context_cache.get(key)
+        if ctx is None:
+            ctx = _make_flow_context(ph, pw, device)
+            flow_context_cache[key] = ctx
+        return ctx
+
+    def inference(img1, img2):
+        """Interpolate a single frame between two input frames."""
+        if img1 is carry["img"]:
+            t1_padded = carry["tensor"]
+            h, w, ph, pw = carry["h"], carry["w"], carry["ph"], carry["pw"]
+        else:
+            t1 = _pil_to_tensor(img1, device)
+            h, w = t1.shape[2], t1.shape[3]
+            ph, pw = _padded_size(h, w)
+            t1_padded = _pad_tensor(t1, ph, pw)
+
+        t2 = _pil_to_tensor(img2, device)
+        t2_padded = _pad_tensor(t2, ph, pw)
+
+        # Carry img2's padded tensor forward in case it's the next pair's img1.
+        carry["img"] = img2
+        carry["tensor"] = t2_padded
+        carry["h"], carry["w"], carry["ph"], carry["pw"] = h, w, ph, pw
+
+        tenFlow_div, backwarp_tenGrid, timestep = get_flow_context(ph, pw)
+
+        with torch.inference_mode():
+            _, _, merged = net(
+                t1_padded,
+                t2_padded,
+                timestep,
+                scale_list,
+                tenFlow_div,
+                backwarp_tenGrid,
+            )
+
+        mid = merged[3][:, :, :h, :w]
+        return float_tensor_to_pil(mid)
+
+    return inference

dw/tasks/model_cache.py

@@ -0,0 +1,39 @@
+import logging
+
+logger = logging.getLogger("dw")
+
+# Loaded task models, keyed by whatever identifies a load - typically
+# (task, model_name, device). step.py runs a task handler once per cartesian
+# product iteration; without this, segmenting 20 images would load the same
+# multi-gigabyte checkpoints 20 times over
+_cache = {}
+
+
+def cached_model(key, factory):
+    """Return the model for key, loading it with factory() on first use.
+
+    Args:
+        key: Hashable identity of the load - include the model name and the
+            device, plus anything else that changes what factory() builds
+        factory: Zero-argument callable performing the actual load
+
+    Returns:
+        The cached or freshly loaded model
+    """
+    if key not in _cache:
+        logger.info(f"Loading task model: {key}")
+        _cache[key] = factory()
+    else:
+        logger.debug(f"Reusing cached task model: {key}")
+    return _cache[key]
+
+
+def clear_model_cache():
+    """Release every cached task model.
+
+    Wired into the worker's memory cleanup - dropping the references here is
+    what lets gc and the allocator actually reclaim the weights.
+    """
+    if _cache:
+        logger.info(f"Clearing {len(_cache)} cached task models")
+    _cache.clear()

dw/tasks/pair_audio.py

@@ -0,0 +1,58 @@
+"""Pair a video with an audio track so the two are saved as one file.
+
+A pipeline that generates its own soundtrack returns the pair together, and the
+result muxes them into a single mp4. Anything that works on the frames alone -
+a latent upsampler, an interpolator, an upscaler - returns frames without it, so
+the soundtrack has to be carried across the step that dropped it. That is what
+this does: it puts the two back together for the step that saves them.
+"""
+
+import logging
+
+from ..result import AudioVideo
+from .audio_utils import as_channels_samples
+
+logger = logging.getLogger("dw")
+
+
+def pair_audio(video, audio, sample_rate=None):
+    """Pair a video's frames with an audio track.
+
+    Args:
+        video: The frames - a frame list, a frame array or tensor, or an
+            AudioVideo whose own soundtrack is replaced by this one
+        audio: The soundtrack - a waveform, or an AudioVideo (or any object
+            carrying '.audio') to take it from, which brings its sample rate
+            along with it
+        sample_rate: Sample rate of the waveform. Required unless `audio`
+            carries one; given here it wins, for a track whose rate was
+            reported wrong
+
+    Returns:
+        One AudioVideo holding the frames and the track
+
+    Raises:
+        ValueError: If no waveform was given, or if no sample rate can be
+            established for the one that was
+    """
+    waveform = getattr(audio, "audio", audio)
+    if waveform is None:
+        raise ValueError(
+            "pair_audio needs an audio track - the video it was given carries none"
+        )
+
+    rate = (
+        sample_rate if sample_rate is not None else getattr(audio, "sample_rate", None)
+    )
+    if rate is None:
+        raise ValueError(
+            "pair_audio needs 'sample_rate' - the audio it was given does not "
+            "carry one of its own"
+        )
+
+    # Frames are left in whatever shape they arrived in - the result saves a frame
+    # list, an array and a tensor alike, and converting a long video here would
+    # cost a copy of the whole thing for nothing
+    frames = video.frames if isinstance(video, AudioVideo) else video
+    logger.debug(f"Pairing frames with audio at {rate} Hz")
+    return AudioVideo(frames, as_channels_samples(waveform), rate)

dw/tasks/qr_code.py

@@ -0,0 +1,19 @@
+import qrcode
+from .image_utils import resize_resample
+
+
+def get_qrcode_image(qr_code_contents, height=768, width=768):
+    # base the resolution off of size - defaulting to 768
+    resolution = max(height, width)
+
+    qr = qrcode.QRCode(
+        version=None,
+        error_correction=qrcode.constants.ERROR_CORRECT_H,
+        box_size=10,
+        border=4,
+    )
+    qr.add_data(qr_code_contents)
+    qr.make(fit=True)
+
+    qrcode_image = qr.make_image(fill_color="black", back_color="white")
+    return resize_resample(qrcode_image, resolution)

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