ndex == int(pipe.num_timestep * 0.4):
                prompt_embeds = callback_kwargs["prompt_embeds"]
                prompt_embeds = prompt_embeds.chunk(2)[-1]

        # update guidance_scale and prompt_embeds
        pipe._guidance_scale = 0.0
        callback_kwargs["prompt_embeds"] = prompt_embeds
        return callback_kwargs Now, you can pass the callback function to the callback_on_step_end parameter and the prompt_embeds to callback_