Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions src/diffusers/models/adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -258,6 +258,12 @@ def __init__(
)

def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
r"""
This function processes the input tensor `x` through the adapter model and returns a list of feature tensors,
each representing information extracted at a different scale from the input.
The length of the list is determined by the number of downsample blocks in the Adapter, as specified
by the `channels` and `num_res_blocks` parameters during initialization.
"""
return self.adapter(x)

@property
Expand Down Expand Up @@ -296,6 +302,12 @@ def __init__(
self.total_downscale_factor = downscale_factor * 2 ** (len(channels) - 1)

def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
r"""
This method processes the input tensor `x` through the FullAdapter model and performs operations including
pixel unshuffling, convolution, and a stack of AdapterBlocks. It returns a list of feature tensors, each capturing information
at a different stage of processing within the FullAdapter model. The number of feature tensors in the list is determined
by the number of downsample blocks specified during initialization.
"""
x = self.unshuffle(x)
x = self.conv_in(x)

Expand Down Expand Up @@ -338,6 +350,10 @@ def __init__(
self.total_downscale_factor = downscale_factor * 2

def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
r"""
This method takes the tensor x as input and processes it through FullAdapterXL model. It consists of operations
including unshuffling pixels, applying convolution layer and appending each block into list of feature tensors.
"""
x = self.unshuffle(x)
x = self.conv_in(x)

Expand Down Expand Up @@ -367,6 +383,11 @@ def __init__(self, in_channels, out_channels, num_res_blocks, down=False):
)

def forward(self, x):
r"""
This method takes tensor x as input and performs operations downsampling and convolutional layers if the
self.downsample and self.in_conv properties of AdapterBlock model are specified. Then it applies a series
of residual blocks to the input tensor.
"""
if self.downsample is not None:
x = self.downsample(x)

Expand All @@ -386,6 +407,10 @@ def __init__(self, channels):
self.block2 = nn.Conv2d(channels, channels, kernel_size=1)

def forward(self, x):
r"""
This method takes input tensor x and applies a convolutional layer, ReLU activation,
and another convolutional layer on the input tensor. It returns addition with the input tensor.
"""
h = x
h = self.block1(h)
h = self.act(h)
Expand Down Expand Up @@ -425,6 +450,10 @@ def __init__(
self.total_downscale_factor = downscale_factor * (2 ** len(channels))

def forward(self, x):
r"""
This method takes the input tensor x and performs downscaling and appends it in list of feature tensors.
Each feature tensor corresponds to a different level of processing within the LightAdapter.
"""
x = self.unshuffle(x)

features = []
Expand All @@ -450,6 +479,10 @@ def __init__(self, in_channels, out_channels, num_res_blocks, down=False):
self.out_conv = nn.Conv2d(mid_channels, out_channels, kernel_size=1)

def forward(self, x):
r"""
This method takes tensor x as input and performs downsampling if required.
Then it applies in convolution layer, a sequence of residual blocks, and out convolutional layer.
"""
if self.downsample is not None:
x = self.downsample(x)

Expand All @@ -468,6 +501,10 @@ def __init__(self, channels):
self.block2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)

def forward(self, x):
r"""
This function takes input tensor x and processes it through one convolutional layer, ReLU activation,
and another convolutional layer and adds it to input tensor.
"""
h = x
h = self.block1(h)
h = self.act(h)
Expand Down