diffusers-workflow 0.4.0__py3-none-any.whl

dw/tasks/tensor_image.py

@@ -0,0 +1,57 @@
+"""
+Shared PIL <-> float tensor conversions.
+
+Consolidates the PIL-to-tensor round trip that was previously hand-rolled
+independently in upscale.py, restore_faces.py, and interpolate_frames.py.
+"""
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def pil_to_float_tensor(image, device, dtype=None):
+    """Convert a PIL image to a (1, 3, H, W) float tensor in [0, 1] on device.
+
+    The image is coerced to RGB first, so single-channel or RGBA inputs are
+    handled consistently. `dtype` defaults to float32 (the numpy source
+    precision); pass e.g. `torch.float16` to cast directly to a model's
+    working precision.
+
+    Args:
+        image: PIL Image
+        device: Target device (str or torch.device)
+        dtype: Optional torch dtype to cast to (default: float32)
+
+    Returns:
+        torch.Tensor of shape (1, 3, H, W)
+    """
+    arr = np.array(image.convert("RGB")).astype(np.float32) / 255.0
+    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(device)
+    if dtype is not None:
+        tensor = tensor.to(dtype=dtype)
+    return tensor
+
+
+def float_tensor_to_pil(tensor):
+    """Convert a (1, 3, H, W) or (3, H, W) float tensor in [0, 1] to a PIL RGB image.
+
+    Quantizes to uint8 via rounding (`.round()`) rather than truncation
+    (a bare truncating cast), matching diffusers' VaeImageProcessor.numpy_to_pil
+    behavior. This matters: with truncation, an exact 8-bit value that
+    round-trips through [0, 1] float (e.g. 128/255) can land a hair below
+    its integer (127.999...) and get chopped down a level instead of
+    landing back on 128. Rounding fixes that, at the cost of at most 1/255
+    of drift per channel versus the old truncating behavior for values that
+    were never exact to begin with.
+
+    Args:
+        tensor: torch.Tensor of shape (1, 3, H, W) or (3, H, W), values in [0, 1]
+
+    Returns:
+        PIL.Image.Image in RGB mode
+    """
+    if tensor.dim() == 4:
+        tensor = tensor.squeeze(0)
+    arr = tensor.permute(1, 2, 0).mul(255).round().clamp(0, 255).byte().cpu().numpy()
+    return Image.fromarray(arr)

dw/tasks/text_generation.py

@@ -0,0 +1,169 @@
+"""
+Text generation via HuggingFace transformers.
+
+Takes a prompt (and optional system prompt) and generates text using a
+local language model. Useful for prompt expansion, rewriting, and
+other text-to-text tasks.
+
+Supplying an image switches to a vision-language model, so the generated
+text can describe what is actually in the picture rather than what the
+prompt guesses is there. That is the difference between an image-conditioned
+workflow whose prompt agrees with its keyframe and one whose prompt fights it.
+"""
+
+import logging
+from transformers import pipeline as hf_pipeline
+from .. import preferred_task_dtype
+from .model_cache import cached_model, hf_pipeline_placement
+
+logger = logging.getLogger("dw")
+
+_DEFAULT_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
+# Small enough to stand in for the old captioning default; override for
+# anything needing real detail
+_DEFAULT_VISION_MODEL = "HuggingFaceTB/SmolVLM-256M-Instruct"
+
+# Greedy decoding against a long, rigid format specification makes the vision
+# models loop - finishing the answer, then repeating its closing sections until
+# they run out of tokens. A penalty stops that without giving up reproducible
+# output, which sampling would. Measured on Qwen3-VL against the H3 prompt
+# spec: 1.05 still looped and filled the whole budget, 1.15 ended on its own at
+# a length matching the format's own guidance. The text models do not need it
+_VISION_REPETITION_PENALTY = 1.15
+
+
+def _image_part(image):
+    """The chat content entry for an image.
+
+    A PIL image goes in as an object; a string is a URL or path the pipeline
+    loads itself, and it has to be declared as such - passing it under "image"
+    would hand the processor a bare string where it expects pixels.
+    """
+    return {"type": "image", "url" if isinstance(image, str) else "image": image}
+
+
+def _build_messages(prompt, system_prompt, image):
+    """Chat messages in the shape the chosen pipeline expects.
+
+    Text-generation models take plain string content. Vision models take a
+    list of typed parts, because the image is a part of the message rather
+    than something alongside it.
+    """
+    messages = []
+    if image is None:
+        if system_prompt is not None:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+    else:
+        if system_prompt is not None:
+            messages.append(
+                {"role": "system", "content": [{"type": "text", "text": system_prompt}]}
+            )
+        messages.append(
+            {
+                "role": "user",
+                "content": [_image_part(image), {"type": "text", "text": prompt}],
+            }
+        )
+    return messages
+
+
+def generate_text(prompt, device="cpu", **kwargs):
+    """Generate text from a prompt using a local language model.
+
+    Args:
+        prompt: The user message / prompt to expand or transform.
+        device: Target device ("cuda", "mps", "cpu").
+        **kwargs:
+            model_name: HuggingFace model ID. Defaults to
+                Qwen/Qwen2.5-1.5B-Instruct, or a small vision-language model
+                when an image is supplied. An image needs a model that can
+                accept one - a text-only model will fail to load as one.
+            system_prompt: Optional system instruction for the model.
+            max_new_tokens: Max tokens to generate (default: 500).
+            image: Optional PIL image, URL or path. Its presence is what
+                selects the vision pipeline.
+            repetition_penalty: Vision pipeline only (default: 1.15). Raise it
+                if a model still repeats itself, or set 1.0 to disable.
+            generate_kwargs: Anything else to hand the model's generate() -
+                no_repeat_ngram_size, top_p, min_new_tokens and so on. Merged
+                over what this function sets, so it can override those too.
+
+    Returns:
+        Generated text string.
+    """
+    # A workflow declaring an optional image passes it through as null when the
+    # caller supplies none, and an empty string is the same statement
+    image = kwargs.get("image", None) or None
+    system_prompt = kwargs.get("system_prompt", None)
+    max_new_tokens = int(kwargs.get("max_new_tokens", 500))
+
+    if image is None:
+        pipeline_task = "text-generation"
+        model_name = kwargs.get("model_name", _DEFAULT_MODEL)
+    else:
+        pipeline_task = "image-text-to-text"
+        model_name = kwargs.get("model_name", _DEFAULT_VISION_MODEL)
+
+    dtype = preferred_task_dtype(device)
+
+    def load_pipe():
+        logger.info(f"Generating text with {model_name} on {device}")
+        placement = hf_pipeline_placement(device)
+        return hf_pipeline(
+            pipeline_task,
+            model=model_name,
+            torch_dtype=dtype,
+            **placement,
+        )
+
+    # The task is part of the identity - the same model name can be loaded
+    # under either pipeline, and they are not interchangeable
+    pipe = cached_model(
+        ("text_generation", pipeline_task, model_name, str(device), str(dtype)),
+        load_pipe,
+    )
+
+    messages = _build_messages(prompt, system_prompt, image)
+
+    # Decoding is greedy, so the same input returns the same text every run -
+    # which is what a workflow wants, and why there is nothing here to seed
+    generation = {"do_sample": False}
+    if image is not None:
+        generation["repetition_penalty"] = float(
+            kwargs.get("repetition_penalty", _VISION_REPETITION_PENALTY)
+        )
+    # Last, so a workflow can override anything decided above
+    generation.update(kwargs.get("generate_kwargs") or {})
+
+    return _generate(pipe, messages, image, max_new_tokens, generation)
+
+
+def _generate(pipe, messages, image, max_new_tokens, generation):
+    """Run the pipeline, which each path calls differently."""
+    if image is None:
+        # This pipeline collects anything it does not name into the arguments it
+        # forwards to generate(), so settings go in as plain keywords
+        results = pipe(
+            messages,
+            max_new_tokens=max_new_tokens,
+            return_full_text=False,
+            **generation,
+        )
+    else:
+        # Images live inside the messages, so the chat goes in as `text` - the
+        # pipeline rejects a chat and an `images` argument together. Generation
+        # settings have to go through generate_kwargs here: anything else this
+        # pipeline does not name explicitly is forwarded to the processor and
+        # dropped, so a bare do_sample=False would leave sampling on. Passing
+        # max_new_tokens both ways is an error, so it stays a direct argument
+        results = pipe(
+            text=messages,
+            max_new_tokens=max_new_tokens,
+            return_full_text=False,
+            generate_kwargs=generation,
+        )
+
+    text = results[0]["generated_text"].strip()
+    logger.info(f"Generated: {text[:100]}{'...' if len(text) > 100 else ''}")
+    return text

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