chore: add V1 flux tests (#742)

* make linter happy

* add tests for qwen-image-edit-2509

* update

* update

* flux schnell test runnable

* update the test score

* make linter happy

* add fp4 results

* fix the test score

* add tests for flux_dev

* update the test score

* add flux.1-krea

* fix the krea tests

* update

* update

* add kontext

* update

* fix kontext

* update

* add flux.1-depth

* add flux-tools

* finish flux tools

* add more flux examples

* update

* update3

* update

* update score

* update

* update

Author: lmxyy

.github/workflows/pr-test.yaml

@@ -92,8 +92,8 @@ jobs:
           source $(conda info --base)/etc/profile.d/conda.sh
           conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
           which python
-          pytest -s -x tests/flux/test_flux_examples.py
-          pytest -s -x tests/v1/test_examples.py
+          pytest -vv -x tests/flux/test_flux_examples.py
+          pytest -vv -x tests/v1/test_examples.py
           python .github/workflows/run_all_tests.py
       - name: clean up
         if: always()

.github/workflows/run_all_tests.py

@@ -1,6 +1,8 @@
 import subprocess
 from pathlib import Path
 
+from tqdm import tqdm
+
 
 def run_all_tests():
     test_dir = Path("tests")
@@ -21,9 +23,9 @@ def run_all_tests():
         print(f"  {test_file}")
 
     failed_tests = []
-    for test_file in test_files:
+    for test_file in tqdm(test_files):
         print(f"Running {test_file} ...")
-        result = subprocess.run(["pytest", "--reruns", "2", "--reruns-delay", "0", "-s", "-x", test_file])
+        result = subprocess.run(["pytest", "--reruns", "2", "--reruns-delay", "0", "-vv", "-x", test_file])
         if result.returncode != 0:
             print(f"Test failed: {test_file}")
             failed_tests.append(test_file)

docker/Dockerfile

@@ -5,8 +5,8 @@
 from nunchaku import NunchakuFluxTransformer2dModel
 from nunchaku.utils import get_precision
 
-image = load_image("https://huggingface.co/mit-han-lab/svdq-int4-flux.1-fill-dev/resolve/main/example.png")
-mask = load_image("https://huggingface.co/mit-han-lab/svdq-int4-flux.1-fill-dev/resolve/main/mask.png")
+image = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup.png")
+mask = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup_mask.png")
 
 precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
 transformer = NunchakuFluxTransformer2dModel.from_pretrained(
@@ -16,7 +16,7 @@
     "black-forest-labs/FLUX.1-Fill-dev", transformer=transformer, torch_dtype=torch.bfloat16
 ).to("cuda")
 image = pipe(
-    prompt="A wooden basket of a cat.",
+    prompt="a white paper cup",
     image=image,
     mask_image=mask,
     height=1024,

docker/Dockerfile.torch27

@@ -0,0 +1,31 @@
+import torch
+from controlnet_aux import CannyDetector
+from diffusers import FluxControlPipeline
+from diffusers.utils import load_image
+
+from nunchaku import NunchakuFluxTransformer2DModelV2
+from nunchaku.utils import get_precision
+
+precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
+transformer = NunchakuFluxTransformer2DModelV2.from_pretrained(
+    f"nunchaku-tech/nunchaku-flux.1-canny-dev/svdq-{precision}_r32-flux.1-canny-dev.safetensors"
+)
+pipe = FluxControlPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Canny-dev", transformer=transformer, torch_dtype=torch.bfloat16
+).to("cuda")
+
+prompt = (
+    "A robot made of exotic candies and chocolates of different kinds. "
+    "The background is filled with confetti and celebratory gifts."
+)
+control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")
+
+processor = CannyDetector()
+control_image = processor(
+    control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024
+)
+
+image = pipe(
+    prompt=prompt, control_image=control_image, height=1024, width=1024, num_inference_steps=20, guidance_scale=30.0
+).images[0]
+image.save(f"flux.1-canny-dev-{precision}.png")

docker/Dockerfile.torch28

@@ -0,0 +1,32 @@
+import torch
+from diffusers import FluxControlPipeline
+from diffusers.utils import load_image
+from image_gen_aux import DepthPreprocessor
+
+from nunchaku import NunchakuFluxTransformer2DModelV2
+from nunchaku.utils import get_precision
+
+precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
+transformer = NunchakuFluxTransformer2DModelV2.from_pretrained(
+    f"nunchaku-tech/nunchaku-flux.1-depth-dev/svdq-{precision}_r32-flux.1-depth-dev.safetensors"
+)
+
+pipe = FluxControlPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Depth-dev",
+    transformer=transformer,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+prompt = (
+    "A robot made of exotic candies and chocolates of different kinds. "
+    "The background is filled with confetti and celebratory gifts."
+)
+control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")
+
+processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
+control_image = processor(control_image)[0].convert("RGB")
+
+image = pipe(
+    prompt=prompt, control_image=control_image, height=1024, width=1024, num_inference_steps=20, guidance_scale=10.0
+).images[0]
+image.save(f"flux.1-depth-dev-{precision}.png")

examples/flux.1-fill-dev.py

@@ -11,5 +11,5 @@
 pipeline = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
 ).to("cuda")
-image = pipeline("A cat holding a sign that says hello world", num_inference_steps=50, guidance_scale=3.5).images[0]
+image = pipeline("A cat holding a sign that says hello world", num_inference_steps=20, guidance_scale=3.5).images[0]
 image.save(f"flux.1-dev-{precision}.png")