From 617c6ffdcba2fc451b89f2c3e03e93eda2ba7ad8 Mon Sep 17 00:00:00 2001 From: ciaranbor Date: Thu, 8 Jan 2026 16:07:46 +0000 Subject: [PATCH] Handle image editing time steps --- src/exo/worker/engines/image/pipeline/runner.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py index 3ee1c794..c8edcf75 100644 --- a/src/exo/worker/engines/image/pipeline/runner.py +++ b/src/exo/worker/engines/image/pipeline/runner.py @@ -229,9 +229,9 @@ class DiffusionRunner: Final denoised latents ready for VAE decoding """ - num_steps = runtime_config.num_inference_steps - runtime_config.init_time_step - - time_steps = tqdm(range(num_steps)) + time_steps = tqdm( + range(runtime_config.init_time_step, runtime_config.num_inference_steps) + ) # Call subscribers for beginning of loop Callbacks.before_loop( @@ -373,7 +373,7 @@ class DiffusionRunner: """ if self.group is None: return self._single_node_step(t, config, latents, prompt_data) - elif t < self.num_sync_steps: + elif t < config.init_time_step + self.num_sync_steps: return self._sync_pipeline( t, config, @@ -528,7 +528,7 @@ class DiffusionRunner: text_seq_len = prompt_embeds.shape[1] hidden_dim = self.adapter.hidden_dim - if t == 0: + if t == config.init_time_step: self._initialize_kv_caches( batch_size=batch_size, num_img_tokens=num_img_tokens, @@ -711,7 +711,10 @@ class DiffusionRunner: start_token, end_token = token_indices[patch_idx] if self.has_joint_blocks: - if not self.is_first_stage or t != self.num_sync_steps: + if ( + not self.is_first_stage + or t != config.init_time_step + self.num_sync_steps + ): if self.is_first_stage: # First stage receives latent-space from last stage (scheduler output) recv_template = patch