diffusers-workflow 0.4.0a3__py3-none-any.whl

dw/tasks/text_sections.py

@@ -0,0 +1,80 @@
+"""
+Reduce a generated block of text to a known set of labelled sections.
+
+A language model asked for a rigid format usually produces it and then keeps
+going - restating the description, appending a summary, or looping until it
+runs out of tokens. Downstream that trailing text is not free: a prompt is
+conditioning, and a pipeline that does not truncate spends memory and attention
+on whatever arrived. Rather than trying to talk the model out of it, keep the
+parts that were asked for and drop the rest.
+"""
+
+import logging
+import re
+
+logger = logging.getLogger("dw")
+
+
+def extract_sections(text, sections, keep_preamble=True):
+    """Keep only the named sections, once each, in the order they are declared.
+
+    A section runs from its `label:` to the end of that paragraph - the format
+    these prompts use puts each field in one continuous paragraph, so a blank
+    line ends it. Anything outside a named section is dropped, which is what
+    removes a trailing restatement whether or not it repeats verbatim.
+
+    Args:
+        text: The generated text.
+        sections: Section labels to keep, in the order they should appear.
+        keep_preamble: Keep any text before the first label. These formats put
+            an instruction line above the fields, which is part of the output.
+
+    Returns:
+        The reassembled text.
+    """
+    if not sections:
+        return text.strip()
+
+    # Matched case-insensitively because the models capitalise labels however
+    # they please - one writes overall_soundscape, another Overall_soundscape,
+    # and dropping a section over its first letter would lose real content. The
+    # declared spelling is what gets written back out
+    labels = "|".join(re.escape(name) for name in sections)
+    matches = list(re.finditer(rf"^({labels}):", text, re.M | re.I))
+    canonical = {name.lower(): name for name in sections}
+
+    if not matches:
+        logger.warning(
+            "No sections of %s found - leaving the text as it is", list(sections)
+        )
+        return text.strip()
+
+    bodies = {}
+    for index, match in enumerate(matches):
+        name = canonical[match.group(1).lower()]
+        if name in bodies:
+            # A repeat is the model starting over, not new content
+            continue
+        # The body ends at the next label or the end of its paragraph, whichever
+        # comes first. A label may sit on its own line, so leading blank space
+        # is skipped before looking for the break that ends it
+        end = matches[index + 1].start() if index + 1 < len(matches) else len(text)
+        body = text[match.end() : end]
+        stripped = body.lstrip()
+        offset = len(body) - len(stripped)
+        paragraph = re.search(r"\n\s*\n", body[offset:])
+        if paragraph:
+            body = body[offset : offset + paragraph.start()]
+        bodies[name] = body.strip()
+
+    kept = [f"{name}: {bodies[name]}" for name in sections if name in bodies]
+
+    preamble = text[: matches[0].start()].strip()
+    if keep_preamble and preamble:
+        kept.insert(0, preamble)
+
+    result = "\n\n".join(kept)
+    dropped = len(text.strip()) - len(result)
+    if dropped > 0:
+        logger.info(f"Trimmed {dropped} characters outside the requested sections")
+    return result

dw/tasks/upscale.py

@@ -0,0 +1,203 @@
+"""
+Image upscaling via spandrel.
+
+Supports 40+ super-resolution architectures (ESRGAN, SwinIR, HAT, DAT, etc.)
+with automatic model detection from .pth/.safetensors files.
+
+Models can be loaded from local files or HuggingFace Hub repos.
+"""
+
+import logging
+import torch
+
+from .model_cache import cached_model
+from .tensor_image import pil_to_float_tensor, float_tensor_to_pil
+
+logger = logging.getLogger("dw")
+
+# Maximum tile size before we switch to tiled processing
+_MAX_PIXELS_NO_TILE = 512 * 512
+
+
+def upscale_image(image, model_name, device="cpu", **kwargs):
+    """Upscale an image using a spandrel-compatible super-resolution model.
+
+    Args:
+        image: PIL Image to upscale
+        model_name: HuggingFace repo ID (e.g., "user/repo") with optional filename,
+                     or local path to .pth/.safetensors file.
+        device: Target device ("cuda", "mps", "cpu")
+        **kwargs:
+            filename: Weight file name within a HF repo (default: auto-detect)
+            tile_size: Tile size for large images (default: 512)
+            tile_overlap: Overlap between tiles in pixels (default: 32)
+
+    Returns:
+        PIL Image (upscaled)
+    """
+    try:
+        from spandrel import ModelLoader, ImageModelDescriptor
+    except ImportError:
+        raise ImportError(
+            "spandrel is required for upscaling. Install with: pip install spandrel"
+        )
+
+    filename = kwargs.get("filename", None)
+    tile_size = kwargs.get("tile_size", 512)
+    tile_overlap = kwargs.get("tile_overlap", 32)
+
+    def load_descriptor():
+        model_path = _resolve_model_path(model_name, filename)
+
+        logger.info(f"Loading upscale model from {model_path}")
+        loader = ModelLoader(device=torch.device(device))
+        result = loader.load_from_file(model_path)
+
+        if not isinstance(result, ImageModelDescriptor):
+            raise ValueError(
+                f"Model is not an image model (got {type(result).__name__}). "
+                f"Only image super-resolution models are supported."
+            )
+
+        logger.info(
+            f"Loaded {result.architecture.name} "
+            f"(scale: {result.scale}x, "
+            f"input: {result.input_channels}ch, "
+            f"output: {result.output_channels}ch)"
+        )
+
+        # Use half precision if supported and on GPU
+        if device != "cpu" and result.supports_half:
+            result.model.half()
+
+        return result
+
+    descriptor = cached_model(
+        ("upscale", model_name, filename, str(device)), load_descriptor
+    )
+    model_dtype = (
+        torch.float16
+        if (device != "cpu" and descriptor.supports_half)
+        else torch.float32
+    )
+
+    # Convert PIL to tensor
+    tensor = pil_to_float_tensor(image, device, dtype=model_dtype)
+
+    h, w = tensor.shape[2], tensor.shape[3]
+
+    if h * w <= _MAX_PIXELS_NO_TILE:
+        logger.debug(f"Upscaling {w}x{h} directly")
+        with torch.inference_mode():
+            output = descriptor(tensor)
+    else:
+        logger.debug(f"Upscaling {w}x{h} with {tile_size}px tiles")
+        output = _tiled_inference(descriptor, tensor, tile_size, tile_overlap)
+
+    # Convert back to PIL
+    result = float_tensor_to_pil(output)
+
+    logger.info(f"Upscaled {w}x{h} -> {result.width}x{result.height}")
+    return result
+
+
+def _resolve_model_path(model_name, filename=None):
+    """Resolve a model name to a local file path.
+
+    Supports:
+        - Local file paths: "/path/to/model.pth"
+        - HuggingFace Hub: "user/repo" (auto-detects .pth/.safetensors)
+        - HuggingFace Hub with filename: model_name="user/repo", filename="model_x4.pth"
+    """
+    import os
+
+    # Local file
+    if os.path.exists(model_name):
+        return model_name
+
+    # HuggingFace Hub
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        raise ImportError(
+            f"Model '{model_name}' is not a local file. "
+            "Install huggingface_hub to download from HF Hub."
+        )
+
+    if filename is not None:
+        logger.debug(f"Downloading {filename} from {model_name}")
+        return hf_hub_download(repo_id=model_name, filename=filename)
+
+    # Auto-detect: list repo files and find a model file
+    from huggingface_hub import list_repo_files
+
+    model_extensions = {".pth", ".pt", ".ckpt", ".safetensors"}
+    try:
+        files = list_repo_files(model_name)
+    except Exception as e:
+        raise ValueError(f"Could not access HuggingFace repo '{model_name}': {e}")
+
+    model_files = [f for f in files if any(f.endswith(ext) for ext in model_extensions)]
+    if not model_files:
+        raise ValueError(
+            f"No model files ({', '.join(model_extensions)}) found in '{model_name}'. "
+            f"Specify 'filename' explicitly."
+        )
+    if len(model_files) > 1:
+        raise ValueError(
+            f"Multiple model files found in '{model_name}': {model_files}. "
+            f"Specify 'filename' explicitly."
+        )
+
+    logger.debug(f"Downloading {model_files[0]} from {model_name}")
+    return hf_hub_download(repo_id=model_name, filename=model_files[0])
+
+
+def _tiled_inference(descriptor, tensor, tile_size, overlap):
+    """Run model inference on overlapping tiles and blend results.
+
+    Splits the input into tiles, runs each through the model, then
+    blends overlapping regions with linear interpolation.
+    """
+    scale = descriptor.scale
+    _, c, h, w = tensor.shape
+    out_h, out_w = h * scale, w * scale
+    output = torch.zeros(1, c, out_h, out_w, device=tensor.device, dtype=tensor.dtype)
+    weight = torch.zeros(1, 1, out_h, out_w, device=tensor.device, dtype=tensor.dtype)
+
+    # Generate tile positions
+    y_positions = list(range(0, h, tile_size - overlap))
+    x_positions = list(range(0, w, tile_size - overlap))
+
+    # Clamp last tile to image boundary
+    y_positions = [min(y, max(0, h - tile_size)) for y in y_positions]
+    x_positions = [min(x, max(0, w - tile_size)) for x in x_positions]
+
+    # Deduplicate
+    y_positions = sorted(set(y_positions))
+    x_positions = sorted(set(x_positions))
+
+    total_tiles = len(y_positions) * len(x_positions)
+    logger.debug(f"Processing {total_tiles} tiles ({tile_size}px, {overlap}px overlap)")
+
+    tile_num = 0
+    for y in y_positions:
+        for x in x_positions:
+            tile_num += 1
+            th = min(tile_size, h - y)
+            tw = min(tile_size, w - x)
+
+            tile = tensor[:, :, y : y + th, x : x + tw]
+
+            with torch.inference_mode():
+                tile_out = descriptor(tile)
+
+            oy, ox = y * scale, x * scale
+            oth, otw = th * scale, tw * scale
+
+            output[:, :, oy : oy + oth, ox : ox + otw] += tile_out
+            weight[:, :, oy : oy + oth, ox : ox + otw] += 1
+
+    # Average overlapping regions
+    output = output / weight.clamp(min=1)
+    return output

dw/tasks/video_utils.py

@@ -0,0 +1,154 @@
+"""Frame access for videos in any of the shapes results carry them.
+
+A video artifact can be a list of PIL images, a numpy array of frames
+(frames, height, width, channels), a torch tensor (frames first, channels
+first or last), or an AudioVideo pairing frames with their generated
+soundtrack. extract_frame gives tasks and the segment-chaining loop one way
+to pull a single frame out of any of them, always as a PIL image.
+"""
+
+import numpy
+import torch
+from PIL import Image
+
+from ..result import AudioVideo
+
+
+def process_video(video, processor, device, kwargs):
+    processor = processor.lower()
+
+    if processor == "get_frame":
+        return get_frame(video, kwargs.get("frame_index", 0))
+
+    if processor == "get_last_frame":
+        return get_frame(video, -1)
+
+    if processor == "get_first_frame":
+        return get_frame(video, 0)
+
+    raise Exception(f"Unknown video processor type: {processor}")
+
+
+def get_frame(video, frame_index=0):
+    return extract_frame(video, frame_index)
+
+
+def extract_frame(video, index):
+    """Pull one frame out of a video, whatever the video's in-memory shape.
+
+    Args:
+        video: List of PIL images, numpy array or torch tensor of frames,
+            an AudioVideo, or a one-video batch wrapping any of those
+        index: Frame to extract; negative indexes count from the end
+
+    Returns:
+        The frame as a PIL image. Frames that already are PIL images are
+        returned as-is, not copied.
+    """
+    return _to_pil(_frames_of(video)[index])
+
+
+def frame_count(video):
+    """Number of frames in a video of any supported shape."""
+    return len(_frames_of(video))
+
+
+def frames_as_pil_list(video):
+    """The video's frames as a list of PIL images.
+
+    Frames that already are PIL images are carried over by identity; array and
+    tensor frames are converted the way extract_frame converts them.
+    """
+    return [_to_pil(frame) for frame in _frames_of(video)]
+
+
+def frames_as_array(video):
+    """The video's frames as one (frames, height, width, channels) uint8 array.
+
+    The shape an argument that takes frames rather than a video wants - LTX-2's
+    keyframe conditions and IC-LoRA references, which the workflow hands what an
+    earlier step generated. One array is also one artifact, where a list of frames
+    would become one artifact per frame and multiply the step that consumed it.
+
+    Frames that already are a channels-last RGB array are converted in a single
+    operation; anything else goes through the same per-frame conversion
+    extract_frame uses.
+    """
+    frames = _frames_of(video)
+
+    if isinstance(frames, numpy.ndarray) and frames.ndim == 4 and frames.shape[-1] == 3:
+        if frames.dtype == numpy.uint8:
+            return frames
+        # Float frames are [0, 1] - diffusers' np output convention
+        return (numpy.clip(frames, 0.0, 1.0) * 255).round().astype(numpy.uint8)
+
+    return numpy.stack(
+        [numpy.asarray(_to_pil(frame).convert("RGB")) for frame in frames]
+    )
+
+
+def _frames_of(video):
+    """Unwrap containers until an indexable run of frames remains."""
+    if isinstance(video, AudioVideo):
+        return _frames_of(video.frames)
+
+    if isinstance(video, list):
+        # A one-video batch - [[frame, ...]] or [ndarray] - unwraps to the video;
+        # a single-frame video - [frame] - is already the frames
+        if len(video) == 1 and not _is_frame(video[0]):
+            return _frames_of(video[0])
+        return video
+
+    if isinstance(video, numpy.ndarray):
+        if video.ndim == 3:  # a lone frame
+            return video[numpy.newaxis, ...]
+        if video.ndim == 5 and video.shape[0] == 1:  # a one-video batch
+            return video[0]
+        return video
+
+    if torch.is_tensor(video):
+        tensor = video.detach().cpu()
+        if tensor.ndim == 5 and tensor.shape[0] == 1:  # a one-video batch
+            tensor = tensor[0]
+        if tensor.ndim == 3:  # a lone frame
+            tensor = tensor.unsqueeze(0)
+        return tensor
+
+    raise TypeError(f"Cannot extract frames from a {type(video).__name__}")
+
+
+def _is_frame(item):
+    """A single image: PIL, or a 3-dim array/tensor (height, width, channels)."""
+    if isinstance(item, Image.Image):
+        return True
+    if isinstance(item, numpy.ndarray) or torch.is_tensor(item):
+        return item.ndim == 3
+    return False
+
+
+def _to_pil(frame):
+    """Convert one frame to a PIL image; PIL frames pass through untouched."""
+    if isinstance(frame, Image.Image):
+        return frame
+
+    if torch.is_tensor(frame):
+        frame = frame.detach().cpu().float().numpy()
+
+    if isinstance(frame, numpy.ndarray):
+        if frame.ndim != 3:
+            raise ValueError(f"A frame must have 3 dimensions, got {frame.ndim}")
+
+        # Channels-first (C, H, W) -> channels-last, the layout PIL expects
+        if frame.shape[0] in (1, 3, 4) and frame.shape[-1] not in (1, 3, 4):
+            frame = numpy.moveaxis(frame, 0, -1)
+
+        if frame.dtype != numpy.uint8:
+            # Float frames are [0, 1] - diffusers' np output convention
+            frame = (numpy.clip(frame, 0.0, 1.0) * 255).round().astype(numpy.uint8)
+
+        if frame.shape[-1] == 1:  # grayscale
+            frame = frame[..., 0]
+
+        return Image.fromarray(frame)
+
+    raise TypeError(f"Cannot convert a {type(frame).__name__} to an image")

dw/tasks/zoe_depth.py

@@ -0,0 +1,71 @@
+import torch
+import matplotlib
+import matplotlib.cm
+import numpy as np
+from PIL import Image
+
+from .model_cache import cached_model
+
+
+def load_zoe(device="cuda"):
+    def load():
+        torch.hub.help(
+            "intel-isl/MiDaS", "DPT_BEiT_L_384"
+        )  # Triggers fresh download of MiDaS repo
+        model_zoe_n = torch.hub.load(
+            "isl-org/ZoeDepth", "ZoeD_NK", pretrained=True
+        ).eval()
+        return model_zoe_n.to(device)
+
+    return cached_model(("zoe_depth", str(device)), load)
+
+
+def colorize(
+    value,
+    vmin=None,
+    vmax=None,
+    cmap="gray_r",
+    invalid_val=-99,
+    invalid_mask=None,
+    background_color=(128, 128, 128, 255),
+    gamma_corrected=False,
+    value_transform=None,
+):
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu().numpy()
+
+    value = value.squeeze()
+    if invalid_mask is None:
+        invalid_mask = value == invalid_val
+    mask = np.logical_not(invalid_mask)
+
+    # normalize
+    vmin = np.percentile(value[mask], 2) if vmin is None else vmin
+    vmax = np.percentile(value[mask], 85) if vmax is None else vmax
+    if vmin != vmax:
+        value = (value - vmin) / (vmax - vmin)  # vmin..vmax
+    else:
+        # Avoid 0-division
+        value = value * 0.0
+
+    # squeeze last dim if it exists
+    # grey out the invalid values
+
+    value[invalid_mask] = np.nan
+    cmapper = matplotlib.cm.get_cmap(cmap)
+    if value_transform:
+        value = value_transform(value)
+        # value = value / value.max()
+    value = cmapper(value, bytes=True)  # (nxmx4)
+
+    # img = value[:, :, :]
+    img = value[...]
+    img[invalid_mask] = background_color
+
+    # gamma correction
+    img = img / 255
+    img = np.power(img, 2.2)
+    img = img * 255
+    img = img.astype(np.uint8)
+    img = Image.fromarray(img)
+    return img