Compare commits

...

3 Commits

Author SHA1 Message Date
Sayak Paul
d500bb94e7 Merge branch 'main' into alert-autofix-2150 2026-03-09 15:57:20 +05:30
Sayak Paul
a08c274c33 [tests] Use tmp_path fixture modular tests (#13194)
* add a test to check modular index consistency

* check for compulsory keys.

* use fixture for tmp_path in modular tests.

* remove unneeded test.

* fix code quality.

* up

* up
2026-03-09 15:17:59 +05:30
Dhruv Nair
ee0d4ed02d Potential fix for code scanning alert no. 2150: Workflow does not contain permissions
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
2026-03-09 10:08:14 +05:30
4 changed files with 55 additions and 61 deletions

View File

@@ -1,5 +1,8 @@
name: Fast GPU Tests on PR
permissions:
contents: read
on:
pull_request:
branches: main

View File

@@ -14,7 +14,6 @@
# limitations under the License.
import random
import tempfile
import numpy as np
import PIL
@@ -129,18 +128,16 @@ class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
return inputs
def test_save_from_pretrained(self):
def test_save_from_pretrained(self, tmp_path):
pipes = []
base_pipe = self.get_pipeline().to(torch_device)
pipes.append(base_pipe)
with tempfile.TemporaryDirectory() as tmpdirname:
base_pipe.save_pretrained(tmpdirname)
pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
pipe.load_components(torch_dtype=torch.float32)
pipe.to(torch_device)
pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
base_pipe.save_pretrained(str(tmp_path))
pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
pipe.load_components(torch_dtype=torch.float32)
pipe.to(torch_device)
pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
pipes.append(pipe)
@@ -212,18 +209,16 @@ class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin):
return inputs
def test_save_from_pretrained(self):
def test_save_from_pretrained(self, tmp_path):
pipes = []
base_pipe = self.get_pipeline().to(torch_device)
pipes.append(base_pipe)
with tempfile.TemporaryDirectory() as tmpdirname:
base_pipe.save_pretrained(tmpdirname)
pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
pipe.load_components(torch_dtype=torch.float32)
pipe.to(torch_device)
pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
base_pipe.save_pretrained(str(tmp_path))
pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
pipe.load_components(torch_dtype=torch.float32)
pipe.to(torch_device)
pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
pipes.append(pipe)

View File

@@ -1,7 +1,6 @@
import gc
import json
import os
import tempfile
from typing import Callable
import pytest
@@ -341,16 +340,15 @@ class ModularPipelineTesterMixin:
assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
def test_save_from_pretrained(self):
def test_save_from_pretrained(self, tmp_path):
pipes = []
base_pipe = self.get_pipeline().to(torch_device)
pipes.append(base_pipe)
with tempfile.TemporaryDirectory() as tmpdirname:
base_pipe.save_pretrained(tmpdirname)
pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
pipe.load_components(torch_dtype=torch.float32)
pipe.to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
pipe.load_components(torch_dtype=torch.float32)
pipe.to(torch_device)
pipes.append(pipe)
@@ -362,32 +360,31 @@ class ModularPipelineTesterMixin:
assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
def test_modular_index_consistency(self):
def test_modular_index_consistency(self, tmp_path):
pipe = self.get_pipeline()
components_spec = pipe._component_specs
components = sorted(components_spec.keys())
with tempfile.TemporaryDirectory() as tmpdir:
pipe.save_pretrained(tmpdir)
index_file = os.path.join(tmpdir, "modular_model_index.json")
assert os.path.exists(index_file)
pipe.save_pretrained(str(tmp_path))
index_file = tmp_path / "modular_model_index.json"
assert index_file.exists()
with open(index_file) as f:
index_contents = json.load(f)
with open(index_file) as f:
index_contents = json.load(f)
compulsory_keys = {"_blocks_class_name", "_class_name", "_diffusers_version"}
for k in compulsory_keys:
assert k in index_contents
compulsory_keys = {"_blocks_class_name", "_class_name", "_diffusers_version"}
for k in compulsory_keys:
assert k in index_contents
to_check_attrs = {"pretrained_model_name_or_path", "revision", "subfolder"}
for component in components:
spec = components_spec[component]
for attr in to_check_attrs:
if getattr(spec, "pretrained_model_name_or_path", None) is not None:
for attr in to_check_attrs:
assert component in index_contents, f"{component} should be present in index but isn't."
attr_value_from_index = index_contents[component][2][attr]
assert getattr(spec, attr) == attr_value_from_index
to_check_attrs = {"pretrained_model_name_or_path", "revision", "subfolder"}
for component in components:
spec = components_spec[component]
for attr in to_check_attrs:
if getattr(spec, "pretrained_model_name_or_path", None) is not None:
for attr in to_check_attrs:
assert component in index_contents, f"{component} should be present in index but isn't."
attr_value_from_index = index_contents[component][2][attr]
assert getattr(spec, attr) == attr_value_from_index
def test_workflow_map(self):
blocks = self.pipeline_blocks_class()
@@ -483,7 +480,7 @@ class TestCustomBlockRequirements:
def test_sequential_block_requirements_save_load(self, tmp_path):
pipe = self.get_dummy_block_pipe()
pipe.save_pretrained(tmp_path)
pipe.save_pretrained(str(tmp_path))
config_path = tmp_path / "modular_config.json"
@@ -508,7 +505,7 @@ class TestCustomBlockRequirements:
logger.setLevel(30)
with CaptureLogger(logger) as cap_logger:
pipe.save_pretrained(tmp_path)
pipe.save_pretrained(str(tmp_path))
template = "{req} was specified in the requirements but wasn't found in the current environment"
msg_xyz = template.format(req="xyz")
@@ -518,7 +515,7 @@ class TestCustomBlockRequirements:
def test_conditional_block_requirements_save_load(self, tmp_path):
pipe = self.get_dummy_conditional_block_pipe()
pipe.save_pretrained(tmp_path)
pipe.save_pretrained(str(tmp_path))
config_path = tmp_path / "modular_config.json"
with open(config_path, "r") as f:
@@ -535,7 +532,7 @@ class TestCustomBlockRequirements:
def test_loop_block_requirements_save_load(self, tmp_path):
pipe = self.get_dummy_loop_block_pipe()
pipe.save_pretrained(tmp_path)
pipe.save_pretrained(str(tmp_path))
config_path = tmp_path / "modular_config.json"
with open(config_path, "r") as f:

View File

@@ -153,25 +153,24 @@ class TestModularCustomBlocks:
output_prompt = output.values["output_prompt"]
assert output_prompt.startswith("Modular diffusers + ")
def test_custom_block_saving_loading(self):
def test_custom_block_saving_loading(self, tmp_path):
custom_block = DummyCustomBlockSimple()
with tempfile.TemporaryDirectory() as tmpdir:
custom_block.save_pretrained(tmpdir)
assert any("modular_config.json" in k for k in os.listdir(tmpdir))
custom_block.save_pretrained(tmp_path)
assert any("modular_config.json" in k for k in os.listdir(tmp_path))
with open(os.path.join(tmpdir, "modular_config.json"), "r") as f:
config = json.load(f)
auto_map = config["auto_map"]
assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"}
with open(os.path.join(tmp_path, "modular_config.json"), "r") as f:
config = json.load(f)
auto_map = config["auto_map"]
assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"}
# For now, the Python script that implements the custom block has to be manually pushed to the Hub.
# This is why, we have to separately save the Python script here.
code_path = os.path.join(tmpdir, "test_modular_pipelines_custom_blocks.py")
with open(code_path, "w") as f:
f.write(CODE_STR)
# For now, the Python script that implements the custom block has to be manually pushed to the Hub.
# This is why, we have to separately save the Python script here.
code_path = os.path.join(tmp_path, "test_modular_pipelines_custom_blocks.py")
with open(code_path, "w") as f:
f.write(CODE_STR)
loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmpdir, trust_remote_code=True)
loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmp_path, trust_remote_code=True)
pipe = loaded_custom_block.init_pipeline()
prompt = "Diffusers is nice"