diffusers-workflow 0.4.0a3__py3-none-any.whl

dw/tasks/audio_utils.py

@@ -0,0 +1,266 @@
+"""Waveform utilities for audio tasks and segment-chained video generation.
+
+Waveforms are handled as (channels, samples) float32 numpy arrays throughout -
+as_channels_samples normalizes the shapes pipelines and files actually produce
+into that layout.
+"""
+
+import io
+import logging
+
+import numpy
+import soundfile
+import torch
+
+from ..security import (
+    validate_path,
+    validate_url,
+    validate_file_extension,
+    ALLOWED_AUDIO_EXTENSIONS,
+)
+
+logger = logging.getLogger("dw")
+
+# A few milliseconds of fade applied on each side of a butt-joined seam so the
+# discontinuity does not click
+DECLICK_MS = 3.0
+
+
+def as_channels_samples(audio):
+    """Normalize a waveform to a (channels, samples) float32 numpy array.
+
+    Accepts torch tensors or numpy arrays shaped (samples,), (channels, samples),
+    (samples, channels), or a one-item batch (1, channels, samples). Channel
+    position is decided the way normalize_audio in result.py decides it: there
+    are always more samples than channels.
+    """
+    if torch.is_tensor(audio):
+        audio = audio.detach().cpu().float().numpy()
+    audio = numpy.asarray(audio, dtype=numpy.float32)
+
+    if audio.ndim == 1:
+        return audio[numpy.newaxis, :]
+
+    if audio.ndim == 3:
+        if audio.shape[0] != 1:
+            raise ValueError(f"Cannot normalize a waveform batch of {audio.shape[0]}")
+        audio = audio[0]
+
+    if audio.ndim != 2:
+        raise ValueError(f"A waveform must have 1-3 dimensions, got {audio.ndim}")
+
+    if audio.shape[0] > audio.shape[1]:  # (samples, channels) -> transpose
+        audio = audio.T
+
+    return numpy.ascontiguousarray(audio)
+
+
+def frames_to_samples(frames, fps, sample_rate):
+    """The number of audio samples spanning a run of video frames."""
+    return int(round(frames / fps * sample_rate))
+
+
+def slice_samples(waveform, start, length):
+    """Cut length samples out of a (channels, samples) waveform from start.
+
+    A slice reaching past the end of the waveform is zero-padded to the
+    requested length, so frame-aligned slicing near the end of a track always
+    yields full-size chunks.
+    """
+    channels, total = waveform.shape
+    piece = waveform[:, start : start + length]
+    if piece.shape[1] < length:
+        padding = numpy.zeros((channels, length - piece.shape[1]), dtype=waveform.dtype)
+        piece = numpy.concatenate([piece, padding], axis=1)
+    return piece
+
+
+def equal_power_crossfade_join(previous, head, following, sample_rate, crossfade_ms):
+    """Join two segments' audio at a seam without changing the total duration.
+
+    previous ends at the seam. head is the audio trimmed off the next segment's
+    start - it covers the same stretch of time as the tail of previous, so the
+    two are blended with an equal-power crossfade over the last
+    min(crossfade_ms, len(head)) of that stretch. following is the next
+    segment's on-timeline audio and is appended unchanged.
+
+    With no head material (nothing was trimmed), the seam gets a short declick
+    ramp instead - a few milliseconds of fade-out and fade-in in place.
+    """
+    previous, head, following = _matched_channels(previous, head, following)
+
+    window = min(
+        int(crossfade_ms / 1000.0 * sample_rate),
+        head.shape[1],
+        previous.shape[1],
+    )
+
+    if window == 0:
+        return _declick_join(previous, following, sample_rate)
+
+    fade_out, fade_in = _equal_power_ramps(window)
+    blended = previous[:, -window:] * fade_out + head[:, -window:] * fade_in
+    return numpy.concatenate([previous[:, :-window], blended, following], axis=1)
+
+
+def crossfade_concat(waveforms, sample_rate, crossfade_ms):
+    """Concatenate waveforms, overlapping each seam by an equal-power crossfade.
+
+    The classic crossfade: each seam overlaps the two waveforms by the fade
+    window, so the result is shorter than the plain sum by one window per seam.
+    """
+    waveforms = [as_channels_samples(waveform) for waveform in waveforms]
+    if not waveforms:
+        raise ValueError("No waveforms to concatenate")
+
+    result = waveforms[0]
+    for following in waveforms[1:]:
+        result, following = _matched_channels(result, following)
+        window = min(
+            int(crossfade_ms / 1000.0 * sample_rate),
+            result.shape[1],
+            following.shape[1],
+        )
+        if window == 0:
+            result = _declick_join(result, following, sample_rate)
+            continue
+
+        fade_out, fade_in = _equal_power_ramps(window)
+        blended = result[:, -window:] * fade_out + following[:, :window] * fade_in
+        result = numpy.concatenate(
+            [result[:, :-window], blended, following[:, window:]], axis=1
+        )
+
+    return result
+
+
+def load_audio(location, base_dir=None):
+    """Load an audio file from a local path or http(s) URL.
+
+    Returns:
+        Tuple of a (channels, samples) float32 waveform and its sample rate
+    """
+    if location.startswith(("http://", "https://")):
+        import requests
+
+        validated_url = validate_url(location)
+        logger.debug(f"Downloading audio from {validated_url}")
+        response = requests.get(validated_url, timeout=60)
+        response.raise_for_status()
+        data, sample_rate = soundfile.read(
+            io.BytesIO(response.content), dtype="float32"
+        )
+    else:
+        validated_path = validate_path(location, base_dir=base_dir, allow_create=False)
+        validate_file_extension(validated_path, ALLOWED_AUDIO_EXTENSIONS)
+        logger.debug(f"Reading audio from {validated_path}")
+        data, sample_rate = soundfile.read(validated_path, dtype="float32")
+
+    # soundfile returns (samples,) or (samples, channels)
+    return as_channels_samples(data), sample_rate
+
+
+def slice_audio(
+    audio,
+    start_seconds=None,
+    duration_seconds=None,
+    start_frame=None,
+    num_frames=None,
+    fps=None,
+    sample_rate=None,
+):
+    """Task command: cut a slice out of an audio track.
+
+    The slice is addressed either in seconds (start_seconds + duration_seconds)
+    or in video frames (start_frame + num_frames + fps). Slices reaching past
+    the end of the track are zero-padded.
+
+    Args:
+        audio: Path or URL of an audio file, or a waveform (which needs
+            sample_rate alongside it)
+        sample_rate: Sample rate of a waveform passed directly; ignored for
+            files, which carry their own
+
+    Returns:
+        The slice as a (samples, channels) float32 array - the layout audio
+        results are saved in
+    """
+    if isinstance(audio, str):
+        waveform, sample_rate = load_audio(audio)
+    else:
+        if sample_rate is None:
+            raise ValueError("slice_audio needs 'sample_rate' with a raw waveform")
+        waveform = as_channels_samples(audio)
+
+    if start_seconds is not None or duration_seconds is not None:
+        if start_seconds is None or duration_seconds is None:
+            raise ValueError(
+                "slice_audio needs both 'start_seconds' and 'duration_seconds'"
+            )
+        start = int(round(start_seconds * sample_rate))
+        length = int(round(duration_seconds * sample_rate))
+    elif start_frame is not None or num_frames is not None:
+        if start_frame is None or num_frames is None or fps is None:
+            raise ValueError(
+                "slice_audio needs 'start_frame', 'num_frames' and 'fps' together"
+            )
+        start = frames_to_samples(start_frame, fps, sample_rate)
+        length = frames_to_samples(num_frames, fps, sample_rate)
+    else:
+        raise ValueError(
+            "slice_audio needs either 'start_seconds'/'duration_seconds' or "
+            "'start_frame'/'num_frames'/'fps'"
+        )
+
+    return slice_samples(waveform, start, length).T
+
+
+def crossfade_audio(audios, crossfade_ms=75, sample_rate=None):
+    """Task command: join audio tracks with an equal-power crossfade.
+
+    Each seam overlaps the two tracks by the fade window, so the result is
+    shorter than the plain sum by one window per seam.
+
+    Args:
+        audios: The waveforms to join, in order
+        crossfade_ms: Length of each crossfade
+        sample_rate: Sample rate of the waveforms
+
+    Returns:
+        The joined track as a (samples, channels) float32 array
+    """
+    if sample_rate is None:
+        raise ValueError("crossfade_audio needs 'sample_rate'")
+    return crossfade_concat(audios, sample_rate, crossfade_ms).T
+
+
+def _equal_power_ramps(window):
+    """Cosine/sine fade curves that sum to constant power across the window."""
+    theta = numpy.linspace(0.0, numpy.pi / 2.0, window, endpoint=False)
+    return numpy.cos(theta, dtype=numpy.float32), numpy.sin(theta, dtype=numpy.float32)
+
+
+def _declick_join(previous, following, sample_rate):
+    """Butt-join two waveforms with a short fade on each side of the seam."""
+    ramp = int(DECLICK_MS / 1000.0 * sample_rate)
+    ramp = min(ramp, previous.shape[1], following.shape[1])
+    if ramp > 0:
+        fade_out, fade_in = _equal_power_ramps(ramp)
+        previous = previous.copy()
+        following = following.copy()
+        previous[:, -ramp:] *= fade_out  # cos: 1 down to ~0
+        following[:, :ramp] *= fade_in  # sin: ~0 up to 1
+    return numpy.concatenate([previous, following], axis=1)
+
+
+def _matched_channels(*waveforms):
+    """Tile mono up so every waveform has the same channel count."""
+    channels = max(waveform.shape[0] for waveform in waveforms)
+    return tuple(
+        (
+            numpy.tile(waveform, (channels, 1))
+            if waveform.shape[0] == 1 and channels > 1
+            else waveform
+        )
+        for waveform in waveforms
+    )

dw/tasks/background_remover.py

@@ -0,0 +1,43 @@
+from PIL import Image
+import torch
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+
+from .model_cache import cached_model
+
+_MODEL_NAME = "briaai/RMBG-2.0"
+
+
+def remove_background(image: Image, device) -> Image:
+    # Model settings
+    def load_model():
+        model = AutoModelForImageSegmentation.from_pretrained(
+            _MODEL_NAME, trust_remote_code=True
+        )
+        model.to(device)
+        model.eval()
+        return model
+
+    model = cached_model(("background_remover", _MODEL_NAME, str(device)), load_model)
+
+    # Data settings
+    transform_image = transforms.Compose(
+        [
+            transforms.Resize((1024, 1024)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+
+    working_copy = image.copy()
+    input_images = transform_image(working_copy).unsqueeze(0).to(device)
+
+    # Prediction
+    with torch.no_grad():
+        preds = model(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(working_copy.size)
+    working_copy.putalpha(mask)
+
+    return working_copy

dw/tasks/borders.py

@@ -0,0 +1,113 @@
+from PIL import Image
+
+
+def add_border_and_mask(
+    image, zoom_all=1.0, zoom_left=0, zoom_right=0, zoom_up=0, zoom_down=0, overlap=0
+):
+    """Adds a black border around the image with individual side control and mask overlap"""
+    orig_width, orig_height = image.size
+
+    # Calculate padding for each side (in pixels)
+    left_pad = int(orig_width * zoom_left)
+    right_pad = int(orig_width * zoom_right)
+    top_pad = int(orig_height * zoom_up)
+    bottom_pad = int(orig_height * zoom_down)
+
+    # Calculate overlap in pixels
+    overlap_left = int(orig_width * overlap)
+    overlap_right = int(orig_width * overlap)
+    overlap_top = int(orig_height * overlap)
+    overlap_bottom = int(orig_height * overlap)
+
+    # If using the all-sides zoom, add it to each side
+    if zoom_all > 1.0:
+        extra_each_side = (zoom_all - 1.0) / 2
+        left_pad += int(orig_width * extra_each_side)
+        right_pad += int(orig_width * extra_each_side)
+        top_pad += int(orig_height * extra_each_side)
+        bottom_pad += int(orig_height * extra_each_side)
+
+    # Calculate new dimensions (ensure they're multiples of 32)
+    new_width = 32 * round((orig_width + left_pad + right_pad) / 32)
+    new_height = 32 * round((orig_height + top_pad + bottom_pad) / 32)
+
+    # Create new image with black border
+    bordered_image = Image.new("RGB", (new_width, new_height), (0, 0, 0))
+    # Paste original image in position
+    paste_x = left_pad
+    paste_y = top_pad
+    bordered_image.paste(image, (paste_x, paste_y))
+
+    # Create mask (white where the border is, black where the original image was)
+    mask = Image.new("L", (new_width, new_height), 255)  # White background
+    # Paste black rectangle with overlap adjustment
+    mask.paste(
+        0,
+        (
+            paste_x + overlap_left,  # Left edge moves right
+            paste_y + overlap_top,  # Top edge moves down
+            paste_x + orig_width - overlap_right,  # Right edge moves left
+            paste_y + orig_height - overlap_bottom,  # Bottom edge moves up
+        ),
+    )
+
+    return {"bordered_image": bordered_image, "mask": mask}
+
+
+def add_border_and_mask_with_size(image, width, height, overlap=0):
+    """
+    Resizes the original image to fit within the target dimensions while maintaining
+    its aspect ratio, then adds borders as needed to reach the exact target size.
+
+    Args:
+        image: PIL Image object
+        width: Target width in pixels
+        height: Target height in pixels
+        overlap: Mask overlap parameter (0-1 range)
+
+    Returns:
+        Dictionary with 'bordered_image' and 'mask'
+    """
+    # Ensure width and height are multiples of 32
+    width = 32 * round(width / 32)
+    height = 32 * round(height / 32)
+
+    # Get original dimensions
+    orig_width, orig_height = image.size
+    orig_aspect = orig_width / orig_height
+    target_aspect = width / height
+
+    # Resize image to fit within target dimensions while maintaining aspect ratio
+    if orig_aspect > target_aspect:
+        # Original is wider than target - fit width
+        new_width = width
+        new_height = int(width / orig_aspect)
+        resized_image = image.resize((new_width, new_height), Image.LANCZOS)
+    else:
+        # Original is taller than target - fit height
+        new_height = height
+        new_width = int(height * orig_aspect)
+        resized_image = image.resize((new_width, new_height), Image.LANCZOS)
+
+    # Now calculate padding to reach target dimensions
+    left_pad = (width - new_width) // 2
+    right_pad = width - new_width - left_pad
+    top_pad = (height - new_height) // 2
+    bottom_pad = height - new_height - top_pad
+
+    # Convert padding to zoom factors (relative to resized dimensions)
+    zoom_left = left_pad / new_width if new_width > 0 else 0
+    zoom_right = right_pad / new_width if new_width > 0 else 0
+    zoom_up = top_pad / new_height if new_height > 0 else 0
+    zoom_down = bottom_pad / new_height if new_height > 0 else 0
+
+    # Call the original function with calculated zoom parameters
+    return add_border_and_mask(
+        resized_image,
+        zoom_all=1.0,
+        zoom_left=zoom_left,
+        zoom_right=zoom_right,
+        zoom_up=zoom_up,
+        zoom_down=zoom_down,
+        overlap=overlap,
+    )

dw/tasks/concat_videos.py

@@ -0,0 +1,80 @@
+"""Concatenate videos - and the audio generated with them - into one video.
+
+The standalone counterpart of what a chained pipeline step does internally:
+frames are joined end to end with an optional head trim on every video after
+the first, and audio tracks are joined at each seam with an equal-power
+crossfade drawn from the trimmed-off material, so video and audio stay in
+sync.
+"""
+
+import logging
+
+from ..result import AudioVideo
+from .audio_utils import (
+    as_channels_samples,
+    equal_power_crossfade_join,
+    frames_to_samples,
+)
+from .video_utils import frames_as_pil_list
+
+logger = logging.getLogger("dw")
+
+
+def concat_videos(videos, trim_frames=0, crossfade_ms=75, fps=None):
+    """Concatenate a list of videos into a single AudioVideo.
+
+    Args:
+        videos: The videos to join, in order - frame lists, frame arrays, or
+            AudioVideos (from gather_videos or previous_result references)
+        trim_frames: Frames dropped from the head of every video after the
+            first - the trim used when each video was generated from the
+            previous one's last frame
+        crossfade_ms: Equal-power crossfade at each audio seam, clamped to
+            the trimmed material
+        fps: Frame rate of the videos - required to join audio when trimming
+
+    Returns:
+        One AudioVideo; its audio is None when no input video carries any
+    """
+    if not isinstance(videos, list) or not videos:
+        raise ValueError("concat_videos needs a non-empty list of videos")
+
+    frames = []
+    audio = None
+    sample_rate = None
+
+    for index, video in enumerate(videos):
+        head_trim = trim_frames if index > 0 else 0
+        frames.extend(frames_as_pil_list(video)[head_trim:])
+
+        if not isinstance(video, AudioVideo) or video.audio is None:
+            continue
+
+        waveform = as_channels_samples(video.audio)
+        if audio is None:
+            audio, sample_rate = waveform, video.sample_rate
+            continue
+
+        if video.sample_rate != sample_rate:
+            raise ValueError(
+                f"Videos carry audio at different sample rates: "
+                f"{sample_rate} then {video.sample_rate}"
+            )
+        if head_trim > 0 and fps is None:
+            raise ValueError(
+                "concat_videos needs 'fps' to trim audio in step with the frames"
+            )
+
+        trim_samples = (
+            frames_to_samples(head_trim, fps, sample_rate) if head_trim else 0
+        )
+        audio = equal_power_crossfade_join(
+            audio,
+            waveform[:, :trim_samples],
+            waveform[:, trim_samples:],
+            sample_rate,
+            crossfade_ms,
+        )
+
+    logger.debug(f"Concatenated {len(videos)} videos into {len(frames)} frames")
+    return AudioVideo(frames, audio, sample_rate)

dw/tasks/depth_estimator.py

@@ -0,0 +1,54 @@
+import torch
+import numpy as np
+from transformers import pipeline
+from torchvision import transforms
+
+from .. import preferred_task_dtype
+from .model_cache import cached_model
+
+
+def make_hint_tensor(image, device, dtype=None):
+    """Estimate depth and return it as a hint tensor for a controlnet pipeline.
+
+    Args:
+        image: Image to estimate depth from
+        device: Device to run the estimator on and place the hint on
+        dtype: Dtype of the hint, defaulting to the one the device works best in.
+            The hint has to match the dtype of the pipeline that consumes it.
+
+    Returns:
+        Depth hint as a tensor of shape (1, 3, height, width)
+    """
+    depth_estimator = cached_model(
+        ("depth_estimator", str(device)),
+        lambda: pipeline("depth-estimation", device=device),
+    )
+
+    image = depth_estimator(image)["depth"]
+    image = np.array(image)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    detected_map = torch.from_numpy(image).float() / 255.0
+    hint = detected_map.permute(2, 0, 1)
+
+    if dtype is None:
+        dtype = preferred_task_dtype(device)
+
+    return hint.unsqueeze(0).to(device=device, dtype=dtype)
+
+
+def make_hint_image(image, device, dtype=None):
+    """Estimate depth and return it as an image.
+
+    Args:
+        image: Image to estimate depth from
+        device: Device to run the estimator on
+        dtype: Dtype to compute the hint in - see make_hint_tensor
+
+    Returns:
+        Depth map as a PIL image
+    """
+    hint = make_hint_tensor(image, device, dtype)
+    # Convert the tensor to a Pillow image
+    to_pil = transforms.ToPILImage()
+    return to_pil(hint[0].float().cpu())

dw/tasks/diffusion_upscale.py

@@ -0,0 +1,109 @@
+"""
+Diffusion-based image upscaling via Stable Diffusion upscale pipelines.
+
+Provides text-guided upscaling with better detail recovery than
+traditional super-resolution models, especially for faces and textures.
+
+Supports two modes:
+  - "x4" (default): StableDiffusionUpscalePipeline (4x, stabilityai/stable-diffusion-x4-upscaler)
+  - "x2": StableDiffusionLatentUpscalePipeline (2x, stabilityai/sd-x2-latent-upscaler)
+"""
+
+import logging
+import torch
+import diffusers
+from .. import preferred_task_dtype
+from .model_cache import cached_model
+
+logger = logging.getLogger("dw")
+
+_MODELS = {
+    "x4": {
+        "pipeline_class": "StableDiffusionUpscalePipeline",
+        "model_name": "stabilityai/stable-diffusion-x4-upscaler",
+    },
+    "x2": {
+        "pipeline_class": "StableDiffusionLatentUpscalePipeline",
+        "model_name": "stabilityai/sd-x2-latent-upscaler",
+    },
+}
+
+
+def diffusion_upscale(image, device="cpu", **kwargs):
+    """Upscale an image using a Stable Diffusion upscale pipeline.
+
+    Args:
+        image: PIL Image to upscale.
+        device: Target device ("cuda", "mps", "cpu").
+        **kwargs:
+            prompt: Text guidance for upscaling (default: "").
+            negative_prompt: Negative text guidance (default: None).
+            mode: "x4" or "x2" (default: "x4").
+            model_name: Override the default model for the selected mode.
+            num_inference_steps: Denoising steps (default: 25).
+            guidance_scale: Classifier-free guidance scale (default: 9.0).
+            noise_level: Noise level for x4 mode (default: 20, ignored for x2).
+
+    Returns:
+        PIL Image (upscaled).
+    """
+    mode = kwargs.get("mode", "x4")
+    if mode not in _MODELS:
+        raise ValueError(f"mode must be one of {sorted(_MODELS.keys())}, got '{mode}'")
+
+    config = _MODELS[mode]
+    model_name = kwargs.get("model_name", config["model_name"])
+    prompt = kwargs.get("prompt", "")
+    negative_prompt = kwargs.get("negative_prompt", None)
+    num_inference_steps = int(kwargs.get("num_inference_steps", 25))
+    guidance_scale = float(kwargs.get("guidance_scale", 9.0))
+    noise_level = int(kwargs.get("noise_level", 20))
+
+    pipeline_class = getattr(diffusers, config["pipeline_class"])
+
+    dtype = preferred_task_dtype(device)
+
+    def load_pipe():
+        logger.info(f"Loading {config['pipeline_class']} from {model_name} to {device}")
+        pipe = pipeline_class.from_pretrained(
+            model_name,
+            torch_dtype=dtype,
+        )
+        pipe.to(device)
+        return pipe
+
+    pipe = cached_model(
+        (
+            "diffusion_upscale",
+            config["pipeline_class"],
+            model_name,
+            str(device),
+            str(dtype),
+        ),
+        load_pipe,
+    )
+
+    call_kwargs = {
+        "prompt": prompt,
+        "image": image,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+    }
+
+    if negative_prompt is not None:
+        call_kwargs["negative_prompt"] = negative_prompt
+
+    if mode == "x4":
+        call_kwargs["noise_level"] = noise_level
+
+    logger.info(
+        f"Upscaling {image.width}x{image.height} with {mode} mode, "
+        f"{num_inference_steps} steps"
+    )
+
+    with torch.inference_mode():
+        result = pipe(**call_kwargs)
+
+    output = result.images[0]
+    logger.info(f"Upscaled to {output.width}x{output.height}")
+    return output

dw/tasks/format_messages.py

@@ -0,0 +1,24 @@
+def format_chat_message(system_prompt, user_message):
+    return {
+        "text_inputs": [
+            {"role": "system", "content": system_prompt},
+            {
+                "role": "user",
+                "content": user_message,
+            },
+        ]
+    }
+
+
+def batch_decode_post_process(processor, task, generated_ids):
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+
+    parsed_answer = processor.post_process_generation(generated_text, task=task)
+
+    return parsed_answer[task]
+
+
+def get_dict_value(dict, key):
+    if key in dict:
+        return dict[key]
+    return None