File tree Expand file tree Collapse file tree 1 file changed +10
-1
lines changed
Expand file tree Collapse file tree 1 file changed +10
-1
lines changed Original file line number Diff line number Diff line change @@ -220,13 +220,22 @@ def enable_free_noise(
220220 self ._free_noise_weighting_scheme = weighting_scheme
221221 self ._free_noise_noise_type = noise_type
222222
223- blocks = [* self .unet .down_blocks , self .unet .mid_block , * self .unet .up_blocks ]
223+ if hasattr (self .unet .mid_block , "motion_modules" ):
224+ blocks = [* self .unet .down_blocks , self .unet .mid_block , * self .unet .up_blocks ]
225+ else :
226+ blocks = [* self .unet .down_blocks , * self .unet .up_blocks ]
227+
224228 for block in blocks :
225229 self ._enable_free_noise_in_block (block )
226230
227231 def disable_free_noise (self ) -> None :
228232 self ._free_noise_context_length = None
229233
234+ if hasattr (self .unet .mid_block , "motion_modules" ):
235+ blocks = [* self .unet .down_blocks , self .unet .mid_block , * self .unet .up_blocks ]
236+ else :
237+ blocks = [* self .unet .down_blocks , * self .unet .up_blocks ]
238+
230239 blocks = [* self .unet .down_blocks , self .unet .mid_block , * self .unet .up_blocks ]
231240 for block in blocks :
232241 self ._disable_free_noise_in_block (block )
You can’t perform that action at this time.
0 commit comments