adding smoke test

Author: amepas

src/maxdiffusion/tests/flux2klein/flux_klein_smoke_test.py

@@ -0,0 +1,104 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import subprocess
+import sys
+import unittest
+import numpy as np
+from PIL import Image
+from skimage.metrics import structural_similarity as ssim
+
+class FluxKleinSmokeTest(unittest.TestCase):
+
+    def setUp(self):
+        self.output_dir = "/tmp/flux_smoke_test_output/"
+        os.makedirs(self.output_dir, exist_ok=True)
+        # Locate the local references directory
+        self.ref_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
+
+    def test_flux_4b_smoke(self):
+        """
+        Runs Flux.2-klein-4B generator end-to-end on TPU VM for a batch of 4 images
+        at 1024x1024 resolution, and validates SSIM consistency against golden references.
+        """
+        cmd = [
+            sys.executable, "src/maxdiffusion/generate_flux2klein.py",
+            "src/maxdiffusion/configs/base_flux2klein.yml",
+            "run_name=flux4b_smoke",
+            "batch_size=4",
+            "height=1024",
+            "width=1024",
+            "use_latents=False",
+            f"output_dir={self.output_dir}",
+            "prompt=A detailed vector illustration of a robotic hummingbird"
+        ]
+        print(f"\n[4B SMOKE] Running command: {' '.join(cmd)}")
+        res = subprocess.run(cmd, capture_output=True, text=True)
+        self.assertEqual(res.returncode, 0, f"4B Generation failed with error:\nSTDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}")
+
+        print("[4B SMOKE] Generation complete. Verifying image SSIM values...")
+        for b in range(4):
+            gen_path = os.path.join(self.output_dir, f"flux2klein_generated_image_b{b}.png")
+            ref_path = os.path.join(self.ref_dir, f"reference_flux4b_b{b}.png")
+            
+            self.assertTrue(os.path.exists(gen_path), f"Generated file not found: {gen_path}")
+            self.assertTrue(os.path.exists(ref_path), f"Reference file not found: {ref_path}")
+            
+            gen_img = np.array(Image.open(gen_path))
+            ref_img = np.array(Image.open(ref_path))
+            
+            val_ssim = ssim(gen_img, ref_img, channel_axis=-1)
+            print(f" -> Batch {b} SSIM vs Reference: {val_ssim:.6f}")
+            self.assertGreater(val_ssim, 0.8, f"SSIM for batch element {b} is below threshold: {val_ssim:.6f}")
+        print("✅ Flux 4B Smoke Test Passed successfully!")
+
+    def test_flux_9b_smoke(self):
+        """
+        Runs Flux.2-klein-9B generator end-to-end on TPU VM for a batch of 4 images
+        at 1024x1024 resolution, and validates SSIM consistency against golden references.
+        """
+        cmd = [
+            sys.executable, "src/maxdiffusion/generate_flux2klein_9B.py",
+            "src/maxdiffusion/configs/base_flux2klein_9B.yml",
+            "run_name=flux9b_smoke",
+            "batch_size=4",
+            "height=1024",
+            "width=1024",
+            "use_latents=False",
+            f"output_dir={self.output_dir}",
+            "prompt=A detailed vector illustration of a robotic hummingbird"
+        ]
+        print(f"\n[9B SMOKE] Running command: {' '.join(cmd)}")
+        res = subprocess.run(cmd, capture_output=True, text=True)
+        self.assertEqual(res.returncode, 0, f"9B Generation failed with error:\nSTDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}")
+
+        print("[9B SMOKE] Generation complete. Verifying image SSIM values...")
+        for b in range(4):
+            gen_path = os.path.join(self.output_dir, f"flux2klein_generated_image_b{b}.png")
+            ref_path = os.path.join(self.ref_dir, f"reference_flux9b_b{b}.png")
+            
+            self.assertTrue(os.path.exists(gen_path), f"Generated file not found: {gen_path}")
+            self.assertTrue(os.path.exists(ref_path), f"Reference file not found: {ref_path}")
+            
+            gen_img = np.array(Image.open(gen_path))
+            ref_img = np.array(Image.open(ref_path))
+            
+            val_ssim = ssim(gen_img, ref_img, channel_axis=-1)
+            print(f" -> Batch {b} SSIM vs Reference: {val_ssim:.6f}")
+            self.assertGreater(val_ssim, 0.8, f"SSIM for batch element {b} is below threshold: {val_ssim:.6f}")
+        print("✅ Flux 9B Smoke Test Passed successfully!")
+
+if __name__ == '__main__':
+    unittest.main()