298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-20 03:13:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove useless code.

Browse Source
This commit is contained in:
comfyanonymous 2023-11-21 17:18:49 -05:00
parent 6a491ebe27
commit 72741105a6
1 changed files with 11 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
comfy/ldm/modules/diffusionmodules/openaimodel.py
View File

@ -28,25 +28,6 @@ class TimestepBlock(nn.Module):
Apply the module to `x` given `emb` timestep embeddings. Apply the module to `x` given `emb` timestep embeddings.
""" """
class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
"""
A sequential module that passes timestep embeddings to the children that
support it as an extra input.
"""
def forward(self, x, emb, context=None, transformer_options={}, output_shape=None):
for layer in self:
if isinstance(layer, TimestepBlock):
x = layer(x, emb)
elif isinstance(layer, SpatialTransformer):
x = layer(x, context, transformer_options)
elif isinstance(layer, Upsample):
x = layer(x, output_shape=output_shape)
else:
x = layer(x)
return x
#This is needed because accelerate makes a copy of transformer_options which breaks "current_index" #This is needed because accelerate makes a copy of transformer_options which breaks "current_index"
def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None): def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None):
for layer in ts: for layer in ts:
@ -54,6 +35,7 @@ def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, out
x = layer(x, emb) x = layer(x, emb)
elif isinstance(layer, SpatialTransformer): elif isinstance(layer, SpatialTransformer):
x = layer(x, context, transformer_options) x = layer(x, context, transformer_options)
if "current_index" in transformer_options:
transformer_options["current_index"] += 1 transformer_options["current_index"] += 1
elif isinstance(layer, Upsample): elif isinstance(layer, Upsample):
x = layer(x, output_shape=output_shape) x = layer(x, output_shape=output_shape)
@ -61,6 +43,15 @@ def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, out
x = layer(x) x = layer(x)
return x return x
class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
"""
A sequential module that passes timestep embeddings to the children that
support it as an extra input.
"""
def forward(self, *args, **kwargs):
return forward_timestep_embed(self, *args, **kwargs)
class Upsample(nn.Module): class Upsample(nn.Module):
""" """
An upsampling layer with an optional convolution. An upsampling layer with an optional convolution.