InvokeAI 6.10.0rc1__py3-none-any.whl → 6.10.0rc2__py3-none-any.whl

invokeai/backend/flux/denoise.py

@@ -1,7 +1,9 @@
+import inspect
 import math
 from typing import Callable
 
 import torch
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from tqdm import tqdm
 
 from invokeai.backend.flux.controlnet.controlnet_flux_output import ControlNetFluxOutput, sum_controlnet_flux_outputs
@@ -35,24 +37,207 @@ def denoise(
     # extra img tokens (sequence-wise) - for Kontext conditioning
     img_cond_seq: torch.Tensor | None = None,
     img_cond_seq_ids: torch.Tensor | None = None,
+    # Optional scheduler for alternative sampling methods
+    scheduler: SchedulerMixin | None = None,
 ):
-    # step 0 is the initial state
-    total_steps = len(timesteps) - 1
-    step_callback(
-        PipelineIntermediateState(
-            step=0,
-            order=1,
-            total_steps=total_steps,
-            timestep=int(timesteps[0]),
-            latents=img,
-        ),
-    )
+    # Determine if we're using a diffusers scheduler or the built-in Euler method
+    use_scheduler = scheduler is not None
+
+    if use_scheduler:
+        # Initialize scheduler with timesteps
+        # The timesteps list contains values in [0, 1] range (sigmas)
+        # LCM should use num_inference_steps (it has its own sigma schedule),
+        # while other schedulers can use custom sigmas if supported
+        is_lcm = scheduler.__class__.__name__ == "FlowMatchLCMScheduler"
+        set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
+        if not is_lcm and "sigmas" in set_timesteps_sig.parameters:
+            # Scheduler supports custom sigmas - use InvokeAI's time-shifted schedule
+            scheduler.set_timesteps(sigmas=timesteps, device=img.device)
+        else:
+            # LCM or scheduler doesn't support custom sigmas - use num_inference_steps
+            # The schedule will be computed by the scheduler itself
+            num_inference_steps = len(timesteps) - 1
+            scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=img.device)
+
+        # For schedulers like Heun, the number of actual steps may differ
+        # (Heun doubles timesteps internally)
+        num_scheduler_steps = len(scheduler.timesteps)
+        # For user-facing step count, use the original number of denoising steps
+        total_steps = len(timesteps) - 1
+    else:
+        total_steps = len(timesteps) - 1
+        num_scheduler_steps = total_steps
+
     # guidance_vec is ignored for schnell.
     guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
 
     # Store original sequence length for slicing predictions
     original_seq_len = img.shape[1]
 
+    # Track the actual step for user-facing progress (accounts for Heun's double steps)
+    user_step = 0
+
+    if use_scheduler:
+        # Use diffusers scheduler for stepping
+        # Use tqdm with total_steps (user-facing steps) not num_scheduler_steps (internal steps)
+        # This ensures progress bar shows 1/8, 2/8, etc. even when scheduler uses more internal steps
+        pbar = tqdm(total=total_steps, desc="Denoising")
+        for step_index in range(num_scheduler_steps):
+            timestep = scheduler.timesteps[step_index]
+            # Convert scheduler timestep (0-1000) to normalized (0-1) for the model
+            t_curr = timestep.item() / scheduler.config.num_train_timesteps
+            t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+
+            # For Heun scheduler, track if we're in first or second order step
+            is_heun = hasattr(scheduler, "state_in_first_order")
+            in_first_order = scheduler.state_in_first_order if is_heun else True
+
+            # Run ControlNet models
+            controlnet_residuals: list[ControlNetFluxOutput] = []
+            for controlnet_extension in controlnet_extensions:
+                controlnet_residuals.append(
+                    controlnet_extension.run_controlnet(
+                        timestep_index=user_step,
+                        total_num_timesteps=total_steps,
+                        img=img,
+                        img_ids=img_ids,
+                        txt=pos_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                        txt_ids=pos_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                        y=pos_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                        timesteps=t_vec,
+                        guidance=guidance_vec,
+                    )
+                )
+
+            merged_controlnet_residuals = sum_controlnet_flux_outputs(controlnet_residuals)
+
+            # Prepare input for model
+            img_input = img
+            img_input_ids = img_ids
+
+            if img_cond is not None:
+                img_input = torch.cat((img_input, img_cond), dim=-1)
+
+            if img_cond_seq is not None:
+                assert img_cond_seq_ids is not None
+                img_input = torch.cat((img_input, img_cond_seq), dim=1)
+                img_input_ids = torch.cat((img_input_ids, img_cond_seq_ids), dim=1)
+
+            pred = model(
+                img=img_input,
+                img_ids=img_input_ids,
+                txt=pos_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                txt_ids=pos_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                y=pos_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                timesteps=t_vec,
+                guidance=guidance_vec,
+                timestep_index=user_step,
+                total_num_timesteps=total_steps,
+                controlnet_double_block_residuals=merged_controlnet_residuals.double_block_residuals,
+                controlnet_single_block_residuals=merged_controlnet_residuals.single_block_residuals,
+                ip_adapter_extensions=pos_ip_adapter_extensions,
+                regional_prompting_extension=pos_regional_prompting_extension,
+            )
+
+            if img_cond_seq is not None:
+                pred = pred[:, :original_seq_len]
+
+            # Get CFG scale for current user step
+            step_cfg_scale = cfg_scale[min(user_step, len(cfg_scale) - 1)]
+
+            if not math.isclose(step_cfg_scale, 1.0):
+                if neg_regional_prompting_extension is None:
+                    raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
+
+                neg_img_input = img
+                neg_img_input_ids = img_ids
+
+                if img_cond is not None:
+                    neg_img_input = torch.cat((neg_img_input, img_cond), dim=-1)
+
+                if img_cond_seq is not None:
+                    neg_img_input = torch.cat((neg_img_input, img_cond_seq), dim=1)
+                    neg_img_input_ids = torch.cat((neg_img_input_ids, img_cond_seq_ids), dim=1)
+
+                neg_pred = model(
+                    img=neg_img_input,
+                    img_ids=neg_img_input_ids,
+                    txt=neg_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                    txt_ids=neg_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                    y=neg_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                    timesteps=t_vec,
+                    guidance=guidance_vec,
+                    timestep_index=user_step,
+                    total_num_timesteps=total_steps,
+                    controlnet_double_block_residuals=None,
+                    controlnet_single_block_residuals=None,
+                    ip_adapter_extensions=neg_ip_adapter_extensions,
+                    regional_prompting_extension=neg_regional_prompting_extension,
+                )
+
+                if img_cond_seq is not None:
+                    neg_pred = neg_pred[:, :original_seq_len]
+                pred = neg_pred + step_cfg_scale * (pred - neg_pred)
+
+            # Use scheduler.step() for the update
+            step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img)
+            img = step_output.prev_sample
+
+            # Get t_prev for inpainting (next sigma value)
+            if step_index + 1 < len(scheduler.sigmas):
+                t_prev = scheduler.sigmas[step_index + 1].item()
+            else:
+                t_prev = 0.0
+
+            if inpaint_extension is not None:
+                img = inpaint_extension.merge_intermediate_latents_with_init_latents(img, t_prev)
+
+            # For Heun, only increment user step after second-order step completes
+            if is_heun:
+                if not in_first_order:
+                    # Second order step completed
+                    user_step += 1
+                    # Only call step_callback if we haven't exceeded total_steps
+                    if user_step <= total_steps:
+                        pbar.update(1)
+                        preview_img = img - t_curr * pred
+                        if inpaint_extension is not None:
+                            preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(
+                                preview_img, 0.0
+                            )
+                        step_callback(
+                            PipelineIntermediateState(
+                                step=user_step,
+                                order=2,
+                                total_steps=total_steps,
+                                timestep=int(t_curr * 1000),
+                                latents=preview_img,
+                            ),
+                        )
+            else:
+                # For LCM and other first-order schedulers
+                user_step += 1
+                # Only call step_callback if we haven't exceeded total_steps
+                # (LCM scheduler may have more internal steps than user-facing steps)
+                if user_step <= total_steps:
+                    pbar.update(1)
+                    preview_img = img - t_curr * pred
+                    if inpaint_extension is not None:
+                        preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_img, 0.0)
+                    step_callback(
+                        PipelineIntermediateState(
+                            step=user_step,
+                            order=1,
+                            total_steps=total_steps,
+                            timestep=int(t_curr * 1000),
+                            latents=preview_img,
+                        ),
+                    )
+
+        pbar.close()
+        return img
+
+    # Original Euler implementation (when scheduler is None)
     for step_index, (t_curr, t_prev) in tqdm(list(enumerate(zip(timesteps[:-1], timesteps[1:], strict=True)))):
         t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
 

invokeai/backend/flux/schedulers.py

@@ -0,0 +1,62 @@
+"""Flow Matching scheduler definitions and mapping.
+
+This module provides the scheduler types and mapping for Flow Matching models
+(Flux and Z-Image), supporting multiple schedulers from the diffusers library.
+"""
+
+from typing import Literal, Type
+
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+    FlowMatchHeunDiscreteScheduler,
+)
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+# Note: FlowMatchLCMScheduler may not be available in all diffusers versions
+try:
+    from diffusers import FlowMatchLCMScheduler
+
+    _HAS_LCM = True
+except ImportError:
+    _HAS_LCM = False
+
+# Scheduler name literal type for type checking
+FLUX_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"]
+
+# Human-readable labels for the UI
+FLUX_SCHEDULER_LABELS: dict[str, str] = {
+    "euler": "Euler",
+    "heun": "Heun (2nd order)",
+    "lcm": "LCM",
+}
+
+# Mapping from scheduler names to scheduler classes
+FLUX_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = {
+    "euler": FlowMatchEulerDiscreteScheduler,
+    "heun": FlowMatchHeunDiscreteScheduler,
+}
+
+if _HAS_LCM:
+    FLUX_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler
+
+
+# Z-Image scheduler types (same schedulers as Flux, both use Flow Matching)
+# Note: Z-Image-Turbo is optimized for ~8 steps with Euler, but other schedulers
+# can be used for experimentation.
+ZIMAGE_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"]
+
+# Human-readable labels for the UI
+ZIMAGE_SCHEDULER_LABELS: dict[str, str] = {
+    "euler": "Euler",
+    "heun": "Heun (2nd order)",
+    "lcm": "LCM",
+}
+
+# Mapping from scheduler names to scheduler classes (same as Flux)
+ZIMAGE_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = {
+    "euler": FlowMatchEulerDiscreteScheduler,
+    "heun": FlowMatchHeunDiscreteScheduler,
+}
+
+if _HAS_LCM:
+    ZIMAGE_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler

invokeai/backend/image_util/pbr_maps/architecture/block.py

@@ -0,0 +1,367 @@
+# Original: https://github.com/joeyballentine/Material-Map-Generator
+# Adopted and optimized for Invoke AI
+
+from collections import OrderedDict
+from typing import Any, List, Literal, Optional
+
+import torch
+import torch.nn as nn
+
+ACTIVATION_LAYER_TYPE = Literal["relu", "leakyrelu", "prelu"]
+NORMALIZATION_LAYER_TYPE = Literal["batch", "instance"]
+PADDING_LAYER_TYPE = Literal["zero", "reflect", "replicate"]
+BLOCK_MODE = Literal["CNA", "NAC", "CNAC"]
+UPCONV_BLOCK_MODE = Literal["nearest", "linear", "bilinear", "bicubic", "trilinear"]
+
+
+def act(act_type: ACTIVATION_LAYER_TYPE, inplace: bool = True, neg_slope: float = 0.2, n_prelu: int = 1):
+    """Helper to select Activation Layer"""
+    if act_type == "relu":
+        layer = nn.ReLU(inplace)
+    elif act_type == "leakyrelu":
+        layer = nn.LeakyReLU(neg_slope, inplace)
+    elif act_type == "prelu":
+        layer = nn.PReLU(num_parameters=n_prelu, init=neg_slope)
+    return layer
+
+
+def norm(norm_type: NORMALIZATION_LAYER_TYPE, nc: int):
+    """Helper to select Normalization Layer"""
+    if norm_type == "batch":
+        layer = nn.BatchNorm2d(nc, affine=True)
+    elif norm_type == "instance":
+        layer = nn.InstanceNorm2d(nc, affine=False)
+    return layer
+
+
+def pad(pad_type: PADDING_LAYER_TYPE, padding: int):
+    """Helper to select Padding Layer"""
+    if padding == 0 or pad_type == "zero":
+        return None
+    if pad_type == "reflect":
+        layer = nn.ReflectionPad2d(padding)
+    elif pad_type == "replicate":
+        layer = nn.ReplicationPad2d(padding)
+    return layer
+
+
+def get_valid_padding(kernel_size: int, dilation: int):
+    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)
+    padding = (kernel_size - 1) // 2
+    return padding
+
+
+def sequential(*args: Any):
+    # Flatten Sequential. It unwraps nn.Sequential.
+    if len(args) == 1:
+        if isinstance(args[0], OrderedDict):
+            raise NotImplementedError("sequential does not support OrderedDict input.")
+        return args[0]  # No sequential is needed.
+    modules: List[nn.Module] = []
+    for module in args:
+        if isinstance(module, nn.Sequential):
+            for submodule in module.children():
+                modules.append(submodule)
+        elif isinstance(module, nn.Module):
+            modules.append(module)
+    return nn.Sequential(*modules)
+
+
+def conv_block(
+    in_nc: int,
+    out_nc: int,
+    kernel_size: int,
+    stride: int = 1,
+    dilation: int = 1,
+    groups: int = 1,
+    bias: bool = True,
+    pad_type: Optional[PADDING_LAYER_TYPE] = "zero",
+    norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+    act_type: Optional[ACTIVATION_LAYER_TYPE] = "relu",
+    mode: BLOCK_MODE = "CNA",
+):
+    """
+    Conv layer with padding, normalization, activation
+    mode: CNA --> Conv -> Norm -> Act
+        NAC --> Norm -> Act --> Conv (Identity Mappings in Deep Residual Networks, ECCV16)
+    """
+    assert mode in ["CNA", "NAC", "CNAC"], f"Wrong conv mode [{mode}]"
+    padding = get_valid_padding(kernel_size, dilation)
+    p = pad(pad_type, padding) if pad_type else None
+    padding = padding if pad_type == "zero" else 0
+
+    c = nn.Conv2d(
+        in_nc,
+        out_nc,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=padding,
+        dilation=dilation,
+        bias=bias,
+        groups=groups,
+    )
+    a = act(act_type) if act_type else None
+    match mode:
+        case "CNA":
+            n = norm(norm_type, out_nc) if norm_type else None
+            return sequential(p, c, n, a)
+        case "NAC":
+            if norm_type is None and act_type is not None:
+                a = act(act_type, inplace=False)
+            n = norm(norm_type, in_nc) if norm_type else None
+            return sequential(n, a, p, c)
+        case "CNAC":
+            n = norm(norm_type, in_nc) if norm_type else None
+            return sequential(n, a, p, c)
+
+
+class ConcatBlock(nn.Module):
+    # Concat the output of a submodule to its input
+    def __init__(self, submodule: nn.Module):
+        super(ConcatBlock, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x: torch.Tensor):
+        output = torch.cat((x, self.sub(x)), dim=1)
+        return output
+
+    def __repr__(self):
+        tmpstr = "Identity .. \n|"
+        modstr = self.sub.__repr__().replace("\n", "\n|")
+        tmpstr = tmpstr + modstr
+        return tmpstr
+
+
+class ShortcutBlock(nn.Module):
+    # Elementwise sum the output of a submodule to its input
+    def __init__(self, submodule: nn.Module):
+        super(ShortcutBlock, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x: torch.Tensor):
+        output = x + self.sub(x)
+        return output
+
+    def __repr__(self):
+        tmpstr = "Identity + \n|"
+        modstr = self.sub.__repr__().replace("\n", "\n|")
+        tmpstr = tmpstr + modstr
+        return tmpstr
+
+
+class ShortcutBlockSPSR(nn.Module):
+    # Elementwise sum the output of a submodule to its input
+    def __init__(self, submodule: nn.Module):
+        super(ShortcutBlockSPSR, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x: torch.Tensor):
+        return x, self.sub
+
+    def __repr__(self):
+        tmpstr = "Identity + \n|"
+        modstr = self.sub.__repr__().replace("\n", "\n|")
+        tmpstr = tmpstr + modstr
+        return tmpstr
+
+
+class ResNetBlock(nn.Module):
+    """
+    ResNet Block, 3-3 style
+    with extra residual scaling used in EDSR
+    (Enhanced Deep Residual Networks for Single Image Super-Resolution, CVPRW 17)
+    """
+
+    def __init__(
+        self,
+        in_nc: int,
+        mid_nc: int,
+        out_nc: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        pad_type: PADDING_LAYER_TYPE = "zero",
+        norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+        act_type: Optional[ACTIVATION_LAYER_TYPE] = "relu",
+        mode: BLOCK_MODE = "CNA",
+        res_scale: int = 1,
+    ):
+        super(ResNetBlock, self).__init__()
+        conv0 = conv_block(
+            in_nc, mid_nc, kernel_size, stride, dilation, groups, bias, pad_type, norm_type, act_type, mode
+        )
+        if mode == "CNA":
+            act_type = None
+        if mode == "CNAC":  # Residual path: |-CNAC-|
+            act_type = None
+            norm_type = None
+        conv1 = conv_block(
+            mid_nc, out_nc, kernel_size, stride, dilation, groups, bias, pad_type, norm_type, act_type, mode
+        )
+
+        self.res = sequential(conv0, conv1)
+        self.res_scale = res_scale
+
+    def forward(self, x: torch.Tensor):
+        res = self.res(x).mul(self.res_scale)
+        return x + res
+
+
+class ResidualDenseBlock_5C(nn.Module):
+    """
+    Residual Dense Block
+    style: 5 convs
+    The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18)
+    """
+
+    def __init__(
+        self,
+        nc: int,
+        kernel_size: int = 3,
+        gc: int = 32,
+        stride: int = 1,
+        bias: bool = True,
+        pad_type: PADDING_LAYER_TYPE = "zero",
+        norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+        act_type: ACTIVATION_LAYER_TYPE = "leakyrelu",
+        mode: BLOCK_MODE = "CNA",
+    ):
+        super(ResidualDenseBlock_5C, self).__init__()
+        # gc: growth channel, i.e. intermediate channels
+        self.conv1 = conv_block(
+            nc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, norm_type=norm_type, act_type=act_type, mode=mode
+        )
+        self.conv2 = conv_block(
+            nc + gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+        )
+        self.conv3 = conv_block(
+            nc + 2 * gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+        )
+        self.conv4 = conv_block(
+            nc + 3 * gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+        )
+        if mode == "CNA":
+            last_act = None
+        else:
+            last_act = act_type
+        self.conv5 = conv_block(
+            nc + 4 * gc, nc, 3, stride, bias=bias, pad_type=pad_type, norm_type=norm_type, act_type=last_act, mode=mode
+        )
+
+    def forward(self, x: torch.Tensor):
+        x1 = self.conv1(x)
+        x2 = self.conv2(torch.cat((x, x1), 1))
+        x3 = self.conv3(torch.cat((x, x1, x2), 1))
+        x4 = self.conv4(torch.cat((x, x1, x2, x3), 1))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5.mul(0.2) + x
+
+
+class RRDB(nn.Module):
+    """
+    Residual in Residual Dense Block
+    (ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks)
+    """
+
+    def __init__(
+        self,
+        nc: int,
+        kernel_size: int = 3,
+        gc: int = 32,
+        stride: int = 1,
+        bias: bool = True,
+        pad_type: PADDING_LAYER_TYPE = "zero",
+        norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+        act_type: ACTIVATION_LAYER_TYPE = "leakyrelu",
+        mode: BLOCK_MODE = "CNA",
+    ):
+        super(RRDB, self).__init__()
+        self.RDB1 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, norm_type, act_type, mode)
+        self.RDB2 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, norm_type, act_type, mode)
+        self.RDB3 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, norm_type, act_type, mode)
+
+    def forward(self, x: torch.Tensor):
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
+        return out.mul(0.2) + x
+
+
+# Upsampler
+def pixelshuffle_block(
+    in_nc: int,
+    out_nc: int,
+    upscale_factor: int = 2,
+    kernel_size: int = 3,
+    stride: int = 1,
+    bias: bool = True,
+    pad_type: PADDING_LAYER_TYPE = "zero",
+    norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+    act_type: ACTIVATION_LAYER_TYPE = "relu",
+):
+    """
+    Pixel shuffle layer
+    (Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional
+    Neural Network, CVPR17)
+    """
+    conv = conv_block(
+        in_nc,
+        out_nc * (upscale_factor**2),
+        kernel_size,
+        stride,
+        bias=bias,
+        pad_type=pad_type,
+        norm_type=None,
+        act_type=None,
+    )
+    pixel_shuffle = nn.PixelShuffle(upscale_factor)
+
+    n = norm(norm_type, out_nc) if norm_type else None
+    a = act(act_type) if act_type else None
+    return sequential(conv, pixel_shuffle, n, a)
+
+
+def upconv_block(
+    in_nc: int,
+    out_nc: int,
+    upscale_factor: int = 2,
+    kernel_size: int = 3,
+    stride: int = 1,
+    bias: bool = True,
+    pad_type: PADDING_LAYER_TYPE = "zero",
+    norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+    act_type: ACTIVATION_LAYER_TYPE = "relu",
+    mode: UPCONV_BLOCK_MODE = "nearest",
+):
+    # Adopted from https://distill.pub/2016/deconv-checkerboard/
+    upsample = nn.Upsample(scale_factor=upscale_factor, mode=mode)
+    conv = conv_block(
+        in_nc, out_nc, kernel_size, stride, bias=bias, pad_type=pad_type, norm_type=norm_type, act_type=act_type
+    )
+    return sequential(upsample, conv)