jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 15:04:50 +08:00
Code Issues Activity

Add some more transformer hooks and move tomesd to comfy_extras.

Browse Source 
Tomesd now uses q instead of x to decide which tokens to merge because
it seems to give better results.
This commit is contained in:
comfyanonymous
2023-06-23 20:17:45 -04:00
parent fa28d7334b
commit 05676942b7
5 changed files with 114 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -830,17 +830,20 @@ class UNetModel(nn.Module):
h = x.type(self.dtype)
for id, module in enumerate(self.input_blocks):
transformer_options["block"] = ("input", id)
h = forward_timestep_embed(module, h, emb, context, transformer_options)
if control is not None and 'input' in control and len(control['input']) > 0:
ctrl = control['input'].pop()
if ctrl is not None:
h += ctrl
hs.append(h)
transformer_options["block"] = ("middle", 0)
h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options)
if control is not None and 'middle' in control and len(control['middle']) > 0:
h += control['middle'].pop()
for module in self.output_blocks:
for id, module in enumerate(self.output_blocks):
transformer_options["block"] = ("output", id)
hsp = hs.pop()
if control is not None and 'output' in control and len(control['output']) > 0:
ctrl = control['output'].pop()