Use standard print flush=True for diagnostic logging

Author: amepas

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -250,7 +250,7 @@ def __call__(
 
       # Stage Sync 1: Phase A Complete
       multihost_utils.sync_global_devices("phase_a_complete")
-      max_logging.log(f"{host_prefix} Passed Phase A Sync Barrier (phase_a_complete) successfully! ✅", flush=True)
+      print(f"{host_prefix} Passed Phase A Sync Barrier (phase_a_complete) successfully! ✅", flush=True)
 
       # Shard pipeline batch inputs across data axis ("data") for SPMD multi-host execution
       data_sharding = jax.sharding.NamedSharding(self.mesh, P("data"))
@@ -267,7 +267,7 @@ def put_data_on_devices(x, sharding):
       txt_ids_val = put_data_on_devices(txt_ids_val, data_sharding)
       img_ids_val = put_data_on_devices(img_ids_val, data_sharding)
 
-      max_logging.log(
+      print(
           f"{host_prefix} DIAGNOSTIC TENSORS BEFORE PHASE B:\n"
           f"  latents_jax: shape={latents_jax.shape}, dtype={latents_jax.dtype}, sharding={getattr(latents_jax, 'sharding', None)}\n"
           f"  prompt_embeds_jax: shape={prompt_embeds_jax.shape}, dtype={prompt_embeds_jax.dtype}, sharding={getattr(prompt_embeds_jax, 'sharding', None)}\n"
@@ -278,12 +278,12 @@ def put_data_on_devices(x, sharding):
 
       # Stage Sync 2: Pre-Phase B Start
       multihost_utils.sync_global_devices("pre_phase_b_start")
-      max_logging.log(f"{host_prefix} Passed Pre-Phase B Sync Barrier (pre_phase_b_start) successfully! ✅", flush=True)
+      print(f"{host_prefix} Passed Pre-Phase B Sync Barrier (pre_phase_b_start) successfully! ✅", flush=True)
 
       # ---------------------------------------------------------------------
       # PHASE B: Denoising Loop (Flux Transformer - Fused JAX fori_loop)
       # ---------------------------------------------------------------------
-      max_logging.log(
+      print(
           f"{host_prefix} [PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images...",
           flush=True,
       )
@@ -326,7 +326,7 @@ def step_fn(step_idx, current_latents):
         )
         latents_jax.block_until_ready()
       except Exception as e:
-        max_logging.log(f"❌ {host_prefix} EXCEPTION IN DENOISE LOOP: {e}", flush=True)
+        print(f"❌ {host_prefix} EXCEPTION IN DENOISE LOOP: {e}", flush=True)
         import traceback
 
         traceback.print_exc()
@@ -335,7 +335,7 @@ def step_fn(step_idx, current_latents):
 
       # Stage Sync 3: Phase B Complete
       multihost_utils.sync_global_devices("phase_b_complete")
-      max_logging.log(f"{host_prefix} Passed Phase B Sync Barrier (phase_b_complete) successfully! ✅", flush=True)
+      print(f"{host_prefix} Passed Phase B Sync Barrier (phase_b_complete) successfully! ✅", flush=True)
 
       trace["denoise_loop"] = time.perf_counter() - t0
       max_logging.log(f" -> [TIMING] Denoising Loop (Flux): {trace['denoise_loop']:.4f} seconds ⏱️")