Align generate_flux2klein.py entry point with reference FLUX/SDXL/WAN architecture

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -97,37 +97,10 @@ def encode_prompt(prompt: str, snapshot_dir: str = None, repo_id: str = "black-f
   return prompt_embeds.cpu().numpy()
 
 
-def main(argv):
+def run(config):
   # Enable shardy partitioner for TPU execution
   jax.config.update("jax_use_shardy_partitioner", True)
 
-  # 1. Load configurations
-  config_path = "src/maxdiffusion/configs/base_flux2klein.yml"
-  custom_overrides = []
-  if len(argv) > 1:
-    if argv[1].endswith(".yml") or argv[1].endswith(".yaml"):
-      config_path = argv[1]
-      if len(argv) > 2:
-        custom_overrides = argv[2:]
-    else:
-      custom_overrides = argv[1:]
-
-  max_logging.log(f"Initializing pyconfig with config: {config_path}")
-  default_args = [
-      None,
-      config_path,
-      "run_name=flux2klein_generation",
-      "output_dir=output/",
-  ]
-  default_args.extend(custom_overrides)
-
-  is_interactive = any(arg and "interactive=True" in arg.replace(" ", "") for arg in default_args)
-  if is_interactive:
-    max_logging.log("ℹ️ Interactive mode detected: overriding use_latents=False for dynamic inputs.")
-    default_args.append("use_latents=False")
-
-  pyconfig.initialize(default_args)
-
   # Import modules after jax.distributed.initialize() has run via pyconfig.initialize()
   from jax.experimental import multihost_utils
   from maxdiffusion.max_utils import device_put_replicated
@@ -138,17 +111,6 @@ def main(argv):
   )
   from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline
 
-  def handle_exception(exc_type, exc_value, exc_traceback):
-    rank = jax.process_index()
-    cnt = jax.process_count()
-    print(f"\n❌ [HOST {rank}/{cnt}] UNCAUGHT FATAL EXCEPTION: {exc_value}", flush=True)
-    import traceback
-    traceback.print_exception(exc_type, exc_value, exc_traceback)
-    sys.stdout.flush()
-
-  sys.excepthook = handle_exception
-
-  config = pyconfig.config
   os.makedirs(config.output_dir, exist_ok=True)
 
   # 2. Setup device mesh
@@ -569,6 +531,40 @@ def _put_leaf(param, sharding):
     multihost_utils.sync_global_devices("flux2klein_generation_complete")
 
 
+def main(argv):
+  def handle_exception(exc_type, exc_value, exc_traceback):
+    rank = jax.process_index()
+    cnt = jax.process_count()
+    print(f"\n❌ [HOST {rank}/{cnt}] UNCAUGHT FATAL EXCEPTION: {exc_value}", flush=True)
+    import traceback
+    traceback.print_exception(exc_type, exc_value, exc_traceback)
+    sys.stdout.flush()
+
+  sys.excepthook = handle_exception
+
+  config_path = "src/maxdiffusion/configs/base_flux2klein.yml"
+  custom_overrides = []
+  if len(argv) > 1:
+    if argv[1].endswith(".yml") or argv[1].endswith(".yaml"):
+      config_path = argv[1]
+      if len(argv) > 2:
+        custom_overrides = argv[2:]
+    else:
+      custom_overrides = argv[1:]
+
+  default_args = [
+      None,
+      config_path,
+      "run_name=flux2klein_generation",
+      "output_dir=output/",
+  ]
+  default_args.extend(custom_overrides)
+
+  pyconfig.initialize(default_args)
+  max_utils.ensure_machinelearning_job_runs(pyconfig.config)
+  run(pyconfig.config)
+
+
 if __name__ == "__main__":
   with transformer_engine_context():
     app.run(main)