Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,332 Bytes
1d117d0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import unittest
import os
from modules.user.pipeline import pipeline
class TestPipeline(unittest.TestCase):
def setUp(self):
self.test_prompt = "a cute cat, high quality, detailed"
self.test_img_path = "../_internal/Flux_00001.png" # Make sure this test image exists
def test_basic_generation_small(self):
pipeline(self.test_prompt, 128, 128, number=1)
# Check if output files exist
def test_basic_generation_medium(self):
pipeline(self.test_prompt, 512, 512, number=1)
def test_basic_generation_large(self):
pipeline(self.test_prompt, 1024, 1024, number=1)
def test_hires_fix(self):
pipeline(self.test_prompt, 512, 512, number=1, hires_fix=True)
def test_adetailer(self):
pipeline(
"a portrait of a person, high quality",
512,
512,
number=1,
adetailer=True
)
def test_enhance_prompt(self):
pipeline(
self.test_prompt,
512,
512,
number=1,
enhance_prompt=True
)
def test_img2img(self):
# Skip if test image doesn't exist
if not os.path.exists(self.test_img_path):
self.skipTest("Test image not found")
pipeline(
self.test_img_path,
512,
512,
number=1,
img2img=True
)
def test_stable_fast(self):
resolutions = [(128, 128), (512, 512), (1024, 1024)]
for w, h in resolutions:
pipeline(
self.test_prompt,
w,
h,
number=1,
stable_fast=True
)
def test_reuse_seed(self):
pipeline(
self.test_prompt,
512,
512,
number=2,
reuse_seed=True
)
def test_flux_mode(self):
resolutions = [(128, 128), (512, 512), (1024, 1024)]
for w, h in resolutions:
pipeline(
self.test_prompt,
w,
h,
number=1,
flux_enabled=True
)
if __name__ == '__main__':
unittest.main() |