techfreakworm commited on
Commit
5ac741c
·
unverified ·
1 Parent(s): 9a5065c

test: l3 gpu smoke (t2i base/turbo + controlnet + upscale)

Browse files

Add pytest.mark.gpu smoke tests for real-weight end-to-end validation.
Also add addopts=-m 'not gpu' to pyproject.toml so CI auto-deselects
these tests; run with `pytest -m gpu` to opt in on a GPU box.

Files changed (2) hide show
  1. pyproject.toml +1 -0
  2. tests/test_smoke_gpu.py +67 -0
pyproject.toml CHANGED
@@ -13,6 +13,7 @@ quote-style = "double"
13
  [tool.pytest.ini_options]
14
  testpaths = ["tests"]
15
  python_files = "test_*.py"
 
16
  markers = [
17
  "gpu: requires a GPU (CUDA or MPS); skipped by default",
18
  ]
 
13
  [tool.pytest.ini_options]
14
  testpaths = ["tests"]
15
  python_files = "test_*.py"
16
+ addopts = "-m 'not gpu'"
17
  markers = [
18
  "gpu: requires a GPU (CUDA or MPS); skipped by default",
19
  ]
tests/test_smoke_gpu.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pytest
2
+
3
+ pytestmark = pytest.mark.gpu
4
+
5
+
6
+ @pytest.fixture(scope="module")
7
+ def real_backend():
8
+ """Build a real backend with real weights. ~30 GB download on first run."""
9
+ import backend
10
+ return backend.ZImageStudioBackend()
11
+
12
+
13
+ def test_t2i_turbo_produces_image(real_backend):
14
+ from PIL import Image
15
+ image, meta = real_backend.generate(
16
+ mode="t2i",
17
+ params=dict(prompt="a red apple on a wooden table",
18
+ negative_prompt="", model="Turbo",
19
+ steps=8, cfg=1.0, width=384, height=384, seed=42,
20
+ lora_path=None, lora_strength=0.0),
21
+ )
22
+ assert isinstance(image, Image.Image)
23
+ assert image.size == (384, 384)
24
+ assert meta["model"] == "Turbo"
25
+
26
+
27
+ def test_t2i_base_produces_image(real_backend):
28
+ from PIL import Image
29
+ image, meta = real_backend.generate(
30
+ mode="t2i",
31
+ params=dict(prompt="a red apple on a wooden table",
32
+ negative_prompt="blurry", model="Base",
33
+ steps=15, cfg=4.0, width=384, height=384, seed=42,
34
+ lora_path=None, lora_strength=0.0),
35
+ )
36
+ assert isinstance(image, Image.Image)
37
+
38
+
39
+ def test_controlnet_produces_image(real_backend):
40
+ from PIL import Image
41
+ import numpy as np
42
+ arr = np.random.randint(0, 255, (384, 384, 3), dtype=np.uint8)
43
+ image, meta = real_backend.generate(
44
+ mode="controlnet",
45
+ params=dict(prompt="a portrait of a person, dramatic light",
46
+ input_image=Image.fromarray(arr),
47
+ preprocessor="Canny", controlnet_scale=1.0,
48
+ steps=9, seed=42, lora_path=None, lora_strength=0.0),
49
+ )
50
+ assert isinstance(image, Image.Image)
51
+
52
+
53
+ def test_upscale_produces_image(real_backend, tmp_path):
54
+ from PIL import Image
55
+ import numpy as np
56
+ from huggingface_hub import hf_hub_download
57
+ arr = np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8)
58
+ image, meta = real_backend.generate(
59
+ mode="upscale",
60
+ params=dict(prompt="masterpiece, 8k",
61
+ input_image=Image.fromarray(arr),
62
+ refine_steps=5, refine_denoise=0.33, seed=42,
63
+ lora_path=None, lora_strength=0.0,
64
+ esrgan_model_path=hf_hub_download("xinntao/Real-ESRGAN",
65
+ "RealESRGAN_x4plus.pth")),
66
+ )
67
+ assert image.size == (512, 512)