huggingface
diff --git a/‎tests/modular_pipelines/flux/test_modular_pipeline_flux.py‎
Lines changed: 2 additions & 1 deletion b/‎tests/modular_pipelines/flux/test_modular_pipeline_flux.py‎
Lines changed: 2 additions & 1 deletion
@@ -29,6 +29,7 @@
 )
 
 from ...testing_utils import floats_tensor, torch_device
+from ..test_components_manager import ModularPipelineOffloadTesterMixin
 from ..test_modular_pipelines_common import ModularPipelineTesterMixin
 
 
@@ -45,7 +46,7 @@
 }
 
 
-class TestFluxModularPipelineFast(ModularPipelineTesterMixin):
+class TestFluxModularPipelineFast(ModularPipelineOffloadTesterMixin, ModularPipelineTesterMixin):
     pipeline_class = FluxModularPipeline
     pipeline_blocks_class = FluxAutoBlocks
     pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-modular"
Original file line number	Diff line number	Diff line change
`@@ -29,6 +29,7 @@`
`29`	`29`	`)`
`30`	`30`
`31`	`31`	`from ...testing_utils import floats_tensor, torch_device`
	`32`	`+from ..test_components_manager import ModularPipelineOffloadTesterMixin`
`32`	`33`	`from ..test_modular_pipelines_common import ModularPipelineTesterMixin`
`33`	`34`
`34`	`35`
`@@ -45,7 +46,7 @@`
`45`	`46`	`}`
`46`	`47`
`47`	`48`
`48`		`-class TestFluxModularPipelineFast(ModularPipelineTesterMixin):`
	`49`	`+class TestFluxModularPipelineFast(ModularPipelineOffloadTesterMixin, ModularPipelineTesterMixin):`
`49`	`50`	`pipeline_class = FluxModularPipeline`
`50`	`51`	`pipeline_blocks_class = FluxAutoBlocks`
`51`	`52`	`pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-modular"`