Clean test paths and organize FLUX.2-Klein image edit reference assets

- Move reference images ref_image_0..3 to tests/images/flux2klein/
- Remove dev-specific hardcoded paths from edit_flux2klein_e2e_test.py
- Use standard /tmp output directories across smoke tests
- Format tests with pyink and ruff

Author: amepas

src/maxdiffusion/tests/edit_flux2klein_e2e_test.py

@@ -54,30 +54,46 @@ def setUp(self):
     jax.config.update("jax_default_matmul_precision", "highest")
     jax.config.update("jax_use_shardy_partitioner", True)
 
-    candidates = [
-        "/mnt/data/hf_cache/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots",
-        "/mnt/hyperdisk_weights/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots",
-        os.path.expanduser("~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots"),
-    ]
-    self.model_dir = None
-    for c in candidates:
-      if os.path.exists(c):
-        snaps = os.listdir(c)
-        if snaps:
-          self.model_dir = os.path.join(c, snaps[0])
-          self.transformer_path = os.path.join(self.model_dir, "transformer")
-          self.vae_path = os.path.join(self.model_dir, "vae", "diffusion_pytorch_model.safetensors")
-          self.text_encoder_path = os.path.join(self.model_dir, "text_encoder")
-          self.tokenizer_path = os.path.join(self.model_dir, "tokenizer")
-          if os.path.exists(self.transformer_path) and os.path.exists(self.vae_path):
-            break
-    self.assertIsNotNone(self.model_dir, "FLUX.2-Klein 4B model directory not found!")
-
-    self.output_dir = "/mnt/data/e2e_parity" if os.path.exists("/mnt/data") else "/tmp/e2e_parity"
+    if "FLUX2_KLEIN_4B_MODEL_PATH" in os.environ:
+      self.model_dir = os.environ["FLUX2_KLEIN_4B_MODEL_PATH"]
+    else:
+      hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+      candidates = [
+          os.path.join(hf_home, "hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots"),
+          os.path.join(hf_home, "hub/models--black-forest-labs--FLUX.2-klein-4b/snapshots"),
+          "/mnt/hyperdisk_weights/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots",
+          "/mnt/data/models/flux2klein-4b",
+      ]
+      self.model_dir = None
+      for c in candidates:
+        if os.path.exists(c):
+          if "snapshots" in c:
+            snaps = os.listdir(c)
+            if snaps:
+              self.model_dir = os.path.join(c, snaps[0])
+          else:
+            self.model_dir = c
+          if self.model_dir:
+            self.transformer_path = os.path.join(self.model_dir, "transformer")
+            self.vae_path = os.path.join(self.model_dir, "vae", "diffusion_pytorch_model.safetensors")
+            self.text_encoder_path = os.path.join(self.model_dir, "text_encoder")
+            self.tokenizer_path = os.path.join(self.model_dir, "tokenizer")
+            if os.path.exists(self.transformer_path) and os.path.exists(self.vae_path):
+              break
+      if self.model_dir is None:
+        self.model_dir = "black-forest-labs/FLUX.2-klein-4B"
+
+    if hasattr(self, "model_dir") and self.model_dir and not hasattr(self, "transformer_path"):
+      self.transformer_path = os.path.join(self.model_dir, "transformer")
+      self.vae_path = os.path.join(self.model_dir, "vae", "diffusion_pytorch_model.safetensors")
+      self.text_encoder_path = os.path.join(self.model_dir, "text_encoder")
+      self.tokenizer_path = os.path.join(self.model_dir, "tokenizer")
+
+    self.output_dir = "/tmp/e2e_parity"
     os.makedirs(self.output_dir, exist_ok=True)
 
     # Resolve reference images
-    ref_dir = "/mnt/data/golden_image_edit_data/ref_images"
+    ref_dir = os.path.join(THIS_DIR, "images", "flux2klein")
     self.ref_images = []
     if os.path.exists(ref_dir):
       for i in range(4):

src/maxdiffusion/tests/generate_flux2klein_smoke_test.py

@@ -41,7 +41,7 @@ def test_flux2klein_4b_smoke(self):
     self.assertTrue(os.path.exists(ref_path), f"Reference image not found: {ref_path}")
     base_image = np.array(Image.open(ref_path)).astype(np.uint8)
 
-    output_dir = "/mnt/data/smoke_test_4b" if os.path.exists("/mnt/data") else "/tmp/smoke_test_4b"
+    output_dir = "/tmp/smoke_test_4b"
     os.makedirs(output_dir, exist_ok=True)
     out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
     if os.path.exists(out_path):
@@ -86,7 +86,7 @@ def test_flux2klein_9b_smoke(self):
     self.assertTrue(os.path.exists(ref_path), f"Reference image not found: {ref_path}")
     base_image = np.array(Image.open(ref_path)).astype(np.uint8)
 
-    output_dir = "/mnt/data/smoke_test_9b" if os.path.exists("/mnt/data") else "/tmp/smoke_test_9b"
+    output_dir = "/tmp/smoke_test_9b"
     os.makedirs(output_dir, exist_ok=True)
     out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
     if os.path.exists(out_path):
@@ -134,7 +134,7 @@ def test_flux2klein_4b_image_edit_smoke(self):
     input_img_path = os.path.join(THIS_DIR, "images", "ref_flux2klein_4b.png")
     self.assertTrue(os.path.exists(input_img_path), f"Input reference image not found: {input_img_path}")
 
-    output_dir = "/mnt/data/smoke_test_image_edit_4b" if os.path.exists("/mnt/data") else "/tmp/smoke_test_image_edit_4b"
+    output_dir = "/tmp/smoke_test_image_edit_4b"
     os.makedirs(output_dir, exist_ok=True)
     out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
     if os.path.exists(out_path):
@@ -183,7 +183,7 @@ def test_flux2klein_9b_image_edit_smoke(self):
     input_img_path = os.path.join(THIS_DIR, "images", "ref_flux2klein_4b.png")
     self.assertTrue(os.path.exists(input_img_path), f"Input reference image not found: {input_img_path}")
 
-    output_dir = "/mnt/data/smoke_test_image_edit_9b" if os.path.exists("/mnt/data") else "/tmp/smoke_test_image_edit_9b"
+    output_dir = "/tmp/smoke_test_image_edit_9b"
     os.makedirs(output_dir, exist_ok=True)
     out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
     if os.path.exists(out_path):