diffusers-workflow 0.4.0__py3-none-any.whl

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/compose_text.py

@@ -0,0 +1,74 @@
+"""Assemble one block of text out of parts written once.
+
+A multi-shot workflow says the same things about its characters in every
+shot: who they are, what they are wearing, what their voice sounds like. The
+engine has no interpolation - a reference is the whole value of an argument,
+never a fragment spliced into one - which is deliberate (docs/WORKFLOW_GUIDE.md,
+'no interpolation'): a `{{name}}` inside a prompt would make every prompt a
+template language nobody declared, and resolution order for a half-substituted
+string is a bottomless pit.
+
+The way out is not interpolation but composition. A part is a whole value -
+a variable, an earlier step's output, a stored prompt - and this task joins
+parts in the order they are given. A character bible is then written once,
+as a variable, and named by every shot that needs it; changing the voice
+changes it everywhere, and nothing has to be hand-copied to stay in step.
+
+Positional, not named: the parts are a list, joined in order. A named form
+('{bible} says {line}') would be the interpolation the engine does not have,
+one layer down.
+"""
+
+import logging
+
+logger = logging.getLogger("dw")
+
+
+def compose_text(parts, separator="\n\n", skip_empty=True):
+    """Task command: join parts into one block of text.
+
+    Args:
+        parts: The parts to join, in order. Each is a whole value - usually
+            a "variable:", "prompt:" or "previous_result:" reference the
+            engine has already resolved. Numbers are written out; None is
+            dropped, so an optional part can be a variable left null
+        separator: What goes between the parts. Defaults to a blank line,
+            the paragraph break the prompt formats are written in
+        skip_empty: Drop parts that are None or empty. With it off, an
+            empty part still contributes its separator
+
+    Returns:
+        The joined text
+
+    Raises:
+        ValueError: If parts is not a list, or holds something that is not
+            text or a number - a dict or an image is a sign a reference
+            resolved to something other than what was meant
+    """
+    if not isinstance(parts, list):
+        raise ValueError(
+            f"compose_text needs a list of parts to join, got {type(parts).__name__}"
+        )
+
+    pieces = []
+    for index, part in enumerate(parts):
+        if part is None:
+            if skip_empty:
+                continue
+            part = ""
+        if isinstance(part, bool) or not isinstance(part, (str, int, float)):
+            raise ValueError(
+                f"compose_text part {index} is a {type(part).__name__} - a part "
+                f"is text (or a number), and anything else means the reference "
+                f"in that position resolved to something other than text"
+            )
+        text = part if isinstance(part, str) else str(part)
+        if skip_empty and not text.strip():
+            continue
+        pieces.append(text)
+
+    if not pieces:
+        raise ValueError("compose_text was given nothing to join")
+
+    logger.debug(f"compose_text: joined {len(pieces)} of {len(parts)} parts")
+    return separator.join(pieces)

dw/tasks/concat_videos.py

@@ -0,0 +1,300 @@
+"""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. Cuts have no trimmed material to fade with, so they can instead let the
+outgoing tail ring on across the seam - see `audio_bleed_ms`.
+"""
+
+import logging
+import os
+
+from ..events import emit_log, emit_warning
+from ..result import AudioVideo
+from ..shots import measured_num_samples, nested_shots, shot_record, trimmed_shots
+from .audio_utils import (
+    as_channels_samples,
+    bleed_join,
+    equal_power_crossfade_join,
+    fit_audio_to_frames,
+    frames_to_samples,
+    match_levels as match_track_levels,
+    resample_waveform,
+    warn_on_level_spread,
+)
+from .video_utils import check_same_frame_size, frames_as_pil_list, load_audio_video
+
+logger = logging.getLogger("dw")
+
+
+def video_names(videos):
+    """A name per video, for an error or a warning that has to say which one.
+
+    A caller passes a path, or a previous step's result; only the path says
+    anything by itself, so the rest are named by position - which is what a
+    six-entry `shots` list needs to be actionable ("24000 then 32000" does
+    not say which entry to fix). By the time this runs, an `asset:`/`output:`
+    reference has already been resolved to its absolute path on this server
+    (#390) - naming a shot by that path leaked server layout onto a consumer
+    surface, so a path is trimmed to its file name, the one part that means
+    anything off this box.
+    """
+    return [
+        os.path.basename(original)
+        if isinstance(original, str)
+        else f"video {index + 1}"
+        for index, original in enumerate(videos)
+    ]
+
+
+def concat_videos(
+    videos,
+    trim_frames=0,
+    crossfade_ms=75,
+    audio_bleed_ms=0,
+    audio_bleed_gain_db=0,
+    seam_fade_ms=None,
+    fps=None,
+    match_levels=None,
+    match_levels_dbfs=None,
+    sample_rate=None,
+):
+    """Concatenate a list of videos into a single AudioVideo.
+
+    Args:
+        videos: The videos to join, in order - frame lists, frame arrays,
+            AudioVideos (from previous_result references), or the path or URL
+            of a video file, which is read with the audio muxed into it. Give
+            each video its own entry: one previous_result reference naming a
+            step that produced several videos fans this step out over them,
+            one concatenation per video, rather than joining them
+        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, drawn from
+            the trimmed material - so it has no effect when trim_frames is 0,
+            which is where every cut-based workflow sits. A hard cut's seam is
+            shaped by audio_bleed_ms or seam_fade_ms instead
+        audio_bleed_ms: How long the outgoing video's tail rings on over the head
+            of the next one, at seams with nothing trimmed to crossfade. For
+            cut-based workflows, where every shot is generated independently and
+            a running laugh track would otherwise butt-join into silence.
+            0 (the default) leaves the seam as a plain declicked join
+        audio_bleed_gain_db: Gain applied to the bled tail before it is added,
+            in dB. 0 (the default) is unchanged, full-scale, matching the
+            outgoing material exactly; negative ducks a tail that would
+            otherwise push the seam over 0 dBFS, or that reads as too present
+            against the incoming shot. Has no effect when audio_bleed_ms is 0
+        seam_fade_ms: Fade applied on each side of a seam that gets neither a
+            crossfade nor a bleed. Defaults to the few milliseconds that keep a
+            butt-join from clicking; raise it to a hundred or so for a graceful
+            hard cut on tonal material, which a bleed would only stutter. It is
+            the wrong tool for a continuous bed such as a laugh track or room
+            tone: a fade only deepens the hole a bleed is there to cover
+        fps: Frame rate of the videos - required to join audio when
+            trimming, and the rate the joined file is written at unless
+            the step's result.fps overrides it
+        match_levels: Even the shots' loudness out before joining -
+            "rms" matches perceived level (the measurement
+            get_gallery_metadata reports as mean_dbfs), "peak" matches the
+            loudest sample. Off by default. Shots generated independently
+            land 10 dB apart routinely, and that jump is the one seam
+            artifact none of the fade controls can hide, because it is not
+            at the seam but either side of it. Left off, a spread wide
+            enough to hear is logged as a warning
+        match_levels_dbfs: The level match_levels moves every shot to -
+            defaults to -1 dBFS for "peak" and -20 dBFS for "rms". A shot
+            that would clip at the target is held at -0.5 dBFS peak instead,
+            reported as a match_levels_held warning, with a per-shot log
+            event naming the hold
+        sample_rate: The rate the joined soundtrack is at. Shots that come
+            from different sources routinely carry different rates - a 24 kHz
+            voice clip paired onto a 32 kHz generation - and unlike a level
+            jump that difference has no editorial meaning, so by default the
+            highest rate among the inputs is chosen and the rest are
+            resampled up to it, with a warning naming which. Give this to pin
+            the target instead (#108)
+
+    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")
+
+    # Named before they are loaded: a path is the only thing that names
+    # itself, and the load below replaces it with what it holds
+    names = video_names(videos)
+    # A shot an earlier run already wrote is loaded here rather than by
+    # gather_videos, which reads frames only and would join it silent
+    videos = [load_audio_video(v) if isinstance(v, str) else v for v in videos]
+    clips = [frames_as_pil_list(v) for v in videos]
+    check_same_frame_size(clips, "concat_videos")
+
+    # Every track up front rather than one at a time: levels are matched
+    # across the whole set, so the last shot's loudness has to be known
+    # before the first one is scaled
+    waveforms = [
+        (
+            as_channels_samples(video.audio)
+            if isinstance(video, AudioVideo) and video.audio is not None
+            else None
+        )
+        for video in videos
+    ]
+    # One rate before anything is joined. Shots assembled from different
+    # sources disagree routinely, and the disagreement carries no meaning -
+    # so it is converted rather than refused, which is what made an agent
+    # invent a resample_audio step by hand (#108)
+    rates = [
+        video.sample_rate
+        for video, waveform in zip(videos, waveforms)
+        if waveform is not None and video.sample_rate
+    ]
+    sample_rate = sample_rate or (max(rates) if rates else None)
+    if rates and len(set(rates)) == 1 and rates[0] != sample_rate:
+        # The inputs agree and the caller pinned another rate: converting
+        # to what was asked for is not a decision made on its behalf (#453)
+        emit_log(
+            f"concat_videos: resampling every track from {rates[0]} Hz to the "
+            f"requested {sample_rate} Hz",
+            command="concat_videos",
+            sample_rate=sample_rate,
+        )
+    elif rates and any(rate != sample_rate for rate in rates):
+        # emit_warning rather than logger.warning, for the reason the level
+        # spread below is emitted: resampling every track is an audio
+        # decision made on the caller's behalf, and a caller reading the job
+        # over the API or MCP sees the warnings list and nothing else - the
+        # conversion landing silently is worse than the loud failure it
+        # replaced (#108)
+        per_video = {
+            name: video.sample_rate
+            for name, video in zip(names, videos)
+            if isinstance(video, AudioVideo) and video.audio is not None
+        }
+        emit_warning(
+            "concat_videos: videos carry audio at different sample rates ("
+            + ", ".join(f"{name}: {rate} Hz" for name, rate in per_video.items())
+            + f") - resampling them all to {sample_rate} Hz. Pass "
+            "'sample_rate' to pin a different target, or resample ahead of "
+            "this step with the 'resample_audio' task.",
+            kind="sample_rate_mismatch",
+            command="concat_videos",
+            sample_rate=sample_rate,
+            sample_rates=per_video,
+        )
+    waveforms = [
+        (
+            waveform
+            if waveform is None
+            or not video.sample_rate
+            or video.sample_rate == sample_rate
+            else resample_waveform(waveform, video.sample_rate, sample_rate)
+        )
+        for video, waveform in zip(videos, waveforms)
+    ]
+
+    if match_levels:
+        waveforms = match_track_levels(waveforms, match_levels, match_levels_dbfs)
+    else:
+        warn_on_level_spread(waveforms)
+
+    frames = []
+    audio = None
+    audio_native_rate = None
+    # Where each video landed, measured on the joined picture and track as
+    # they grow - never derived from the frame count, so a track that runs
+    # long shows up here as the samples it actually took (#378)
+    shots = []
+
+    for index, (video, clip) in enumerate(zip(videos, clips)):
+        head_trim = trim_frames if index > 0 else 0
+        start_frame = len(frames)
+        start_sample = audio.shape[1] if audio is not None else 0
+        frames.extend(clip[head_trim:])
+        inner = getattr(video, "shots", None)
+        if inner:
+            video_shots = nested_shots(
+                trimmed_shots(inner, head_trim),
+                start_frame,
+                start_sample if waveforms[index] is not None else None,
+                getattr(video, "sample_rate", None),
+                sample_rate,
+            )
+        else:
+            video_shots = [
+                shot_record(
+                    names[index], start_frame, len(frames) - start_frame, start_sample
+                )
+            ]
+        # Which input this shot came from - named_shots (dw/shots.py) uses
+        # it to place a step's override name on the right shot once an
+        # earlier input has nested more than one of its own (#432)
+        for shot in video_shots:
+            shot["source_index"] = index
+        shots.extend(video_shots)
+
+        if waveforms[index] is None:
+            continue
+
+        waveform = waveforms[index]
+        if audio is None:
+            audio = waveform
+            audio_native_rate = video.sample_rate
+            continue
+
+        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
+        )
+        if trim_samples == 0 and audio_bleed_ms:
+            audio = bleed_join(
+                audio,
+                waveform,
+                sample_rate,
+                audio_bleed_ms,
+                seam_fade_ms,
+                audio_bleed_gain_db,
+                native_sample_rate=audio_native_rate,
+            )
+        else:
+            audio = equal_power_crossfade_join(
+                audio,
+                waveform[:, :trim_samples],
+                waveform[:, trim_samples:],
+                sample_rate,
+                crossfade_ms,
+                seam_fade_ms,
+            )
+        audio_native_rate = video.sample_rate
+
+    # The rate the caller declared, else the rate the first input carries -
+    # either beats the result's 8 fps default (#84)
+    written_fps = fps or next(
+        (v.fps for v in videos if getattr(v, "fps", None)),
+        None,
+    )
+    # Reconciled against the frame grid before shots are measured (#435), so
+    # an input already short of its own grid does not carry its shortfall
+    # into this join's shot map and compound in a later one
+    audio = fit_audio_to_frames(
+        audio, sample_rate, len(frames), written_fps, "concat_videos"
+    )
+
+    # A seam's crossfade leaves the samples before it where they were, so a
+    # shot's track is everything up to where the next measured one began
+    measured_num_samples(shots, _length(audio) if audio is not None else None)
+
+    logger.debug(f"Concatenated {len(videos)} videos into {len(frames)} frames")
+    return AudioVideo(frames, audio, sample_rate, fps=written_fps, shots=shots)
+
+
+def _length(audio):
+    """How many samples a joined track holds, 0 for none."""
+    return 0 if audio is None else audio.shape[1]

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