question about kv_cache

[hrz2000]

OmniGen/OmniGen/scheduler.py

Line 106 in 8d1647c

return self.key_cache[layer_idx], self.value_cache[layer_idx]

Why are timestep and image_latents (collectively referred to as x) not kept during recording kv_cache, but only condition c? This makes the first time step of inference not perform self-attention of x on x, but only cross-attention of x on c.

[staoxiao]

Hi, @hrz2000 , thank you for pointing out this issue. It appears to be a bug in the inference process, and I will carefully check this part of the code later.