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(
220
220
self ._free_noise_weighting_scheme = weighting_scheme
221
221
self ._free_noise_noise_type = noise_type
222
222
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
+
224
228
for block in blocks :
225
229
self ._enable_free_noise_in_block (block )
226
230
227
231
def disable_free_noise (self ) -> None :
228
232
self ._free_noise_context_length = None
229
233
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
+
230
239
blocks = [* self .unet .down_blocks , self .unet .mid_block , * self .unet .up_blocks ]
231
240
for block in blocks :
232
241
self ._disable_free_noise_in_block (block )
You can’t perform that action at this time.
0 commit comments