update a pytest to test all examples

tests/flux/test_flux_examples.py

+import os
+import subprocess
+
+import pytest
+
+EXAMPLES_DIR = "./examples"
+
+example_scripts = [f for f in os.listdir(EXAMPLES_DIR) if f.endswith(".py") and f.startswith("flux")]
+
+
+@pytest.mark.parametrize("script_name", example_scripts)
+def test_example_script_runs(script_name):
+    script_path = os.path.join(EXAMPLES_DIR, script_name)
+    result = subprocess.run(["python", script_path], capture_output=True, text=True)
+    print(f"Running {script_path} -> Return code: {result.returncode}")
+    print(result.stdout)
+    print(result.stderr)
+
+    assert result.returncode == 0, f"{script_path} failed with code {result.returncode}"

tests/sana/test_examples.py

+import os
+import subprocess
+
+import pytest
+
+from nunchaku.utils import get_precision, is_turing
+
+EXAMPLES_DIR = "./examples"
+
+example_scripts = [f for f in os.listdir(EXAMPLES_DIR) if f.endswith(".py") and f.startswith("sana")]
+
+
+@pytest.mark.skipif(
+    is_turing() or get_precision() == "fp4", reason="SANA does not support Turing GPUs or FP4 precision"
+)
+@pytest.mark.parametrize("script_name", example_scripts)
+def test_example_script_runs(script_name):
+    script_path = os.path.join(EXAMPLES_DIR, script_name)
+    result = subprocess.run(["python", script_path], capture_output=True, text=True)
+    print(f"Running {script_path} -> Return code: {result.returncode}")
+    print(result.stdout)
+    print(result.stderr)
+
+    assert result.returncode == 0, f"{script_path} failed with code {result.returncode}"

tests/sana/test_t2i.py

-import pytest
-import torch
-from diffusers import SanaPAGPipeline, SanaPipeline
-
-from nunchaku import NunchakuSanaTransformer2DModel
-from nunchaku.utils import get_precision, is_turing
-
-
-@pytest.mark.skipif(is_turing() or get_precision() == "fp4", reason="Skip tests due to Turing GPUs")
-def test_sana():
-    transformer = NunchakuSanaTransformer2DModel.from_pretrained("mit-han-lab/svdq-int4-sana-1600m")
-    pipe = SanaPipeline.from_pretrained(
-        "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
-        transformer=transformer,
-        variant="bf16",
-        torch_dtype=torch.bfloat16,
-    ).to("cuda")
-    pipe.vae.to(torch.bfloat16)
-    pipe.text_encoder.to(torch.bfloat16)
-
-    prompt = "A cute 🐼 eating 🎋, ink drawing style"
-    image = pipe(
-        prompt=prompt,
-        height=1024,
-        width=1024,
-        guidance_scale=4.5,
-        num_inference_steps=20,
-        generator=torch.Generator().manual_seed(42),
-    ).images[0]
-
-    image.save("sana_1600m.png")
-
-
-@pytest.mark.skipif(is_turing() or get_precision() == "fp4", reason="Skip tests due to Turing GPUs")
-def test_sana_pag():
-    transformer = NunchakuSanaTransformer2DModel.from_pretrained("mit-han-lab/svdq-int4-sana-1600m", pag_layers=8)
-    pipe = SanaPAGPipeline.from_pretrained(
-        "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
-        transformer=transformer,
-        variant="bf16",
-        torch_dtype=torch.bfloat16,
-        pag_applied_layers="transformer_blocks.8",
-    ).to("cuda")
-    pipe._set_pag_attn_processor = lambda *args, **kwargs: None
-
-    pipe.text_encoder.to(torch.bfloat16)
-    pipe.vae.to(torch.bfloat16)
-
-    image = pipe(
-        prompt="A cute 🐼 eating 🎋, ink drawing style",
-        height=1024,
-        width=1024,
-        guidance_scale=5.0,
-        pag_scale=2.0,
-        num_inference_steps=20,
-    ).images[0]
-    image.save("sana_1600m_pag.png")