wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-16 21:42:32 +08:00
Code Issues Actions 16 Packages Projects Releases Wiki Activity

Add support for small flux.2 decoder

Browse Source
This commit is contained in:
huemin-art 2026-04-07 01:46:45 +00:00
parent 61c2387436
commit a0a33ca618
2 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/ldm/models/autoencoder.py
View File

@ -155,6 +155,7 @@ class AutoencodingEngineLegacy(AutoencodingEngine):
def __init__(self, embed_dim: int, **kwargs): def __init__(self, embed_dim: int, **kwargs):
self.max_batch_size = kwargs.pop("max_batch_size", None) self.max_batch_size = kwargs.pop("max_batch_size", None)
ddconfig = kwargs.pop("ddconfig") ddconfig = kwargs.pop("ddconfig")
decoder_ddconfig = kwargs.pop("decoder_ddconfig", ddconfig)
super().__init__( super().__init__(
encoder_config={ encoder_config={
"target": "comfy.ldm.modules.diffusionmodules.model.Encoder", "target": "comfy.ldm.modules.diffusionmodules.model.Encoder",
@ -162,7 +163,7 @@ class AutoencodingEngineLegacy(AutoencodingEngine):
}, },
decoder_config={ decoder_config={
"target": "comfy.ldm.modules.diffusionmodules.model.Decoder", "target": "comfy.ldm.modules.diffusionmodules.model.Decoder",
"params": ddconfig, "params": decoder_ddconfig,
}, },
**kwargs, **kwargs,
) )

11
comfy/sd.py
View File

@ -556,12 +556,19 @@ class VAE:
old_memory_used_decode = self.memory_used_decode old_memory_used_decode = self.memory_used_decode
self.memory_used_decode = lambda shape, dtype: old_memory_used_decode(shape, dtype) * 4.0 self.memory_used_decode = lambda shape, dtype: old_memory_used_decode(shape, dtype) * 4.0
decoder_ch = sd['decoder.conv_in.weight'].shape[0] // ddconfig['ch_mult'][-1]
if decoder_ch != ddconfig['ch']:
decoder_ddconfig = ddconfig.copy()
decoder_ddconfig['ch'] = decoder_ch
else:
decoder_ddconfig = None
if 'post_quant_conv.weight' in sd: if 'post_quant_conv.weight' in sd:
self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1]) self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1], **({"decoder_ddconfig": decoder_ddconfig} if decoder_ddconfig is not None else {}))
else: else:
self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"}, self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig}, encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig},
decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig}) decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': decoder_ddconfig if decoder_ddconfig is not None else ddconfig})
elif "decoder.layers.1.layers.0.beta" in sd: elif "decoder.layers.1.layers.0.beta" in sd:
config = {} config = {}
param_key = None param_key = None