qwen: Implement transformer block prefetching

Author: rattus128

comfy/ldm/qwen_image/model.py

@@ -385,6 +385,9 @@ def _forward(
         hidden_states, img_ids, orig_shape = self.process_img(x)
         num_embeds = hidden_states.shape[1]
 
+        prefetch_queue = comfy.ops.make_prefetch_queue(list(self.transformer_blocks))
+        comfy.ops.prefetch_queue_pop(prefetch_queue, x.device, None)
+
         if ref_latents is not None:
             h = 0
             w = 0
@@ -434,6 +437,7 @@ def _forward(
         blocks_replace = patches_replace.get("dit", {})
 
         for i, block in enumerate(self.transformer_blocks):
+            comfy.ops.prefetch_queue_pop(prefetch_queue, x.device, block)
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -465,6 +469,8 @@ def block_wrap(args):
                     if add is not None:
                         hidden_states[:, :add.shape[1]] += add
 
+        comfy.ops.prefetch_queue_pop(prefetch_queue, x.device, block)
+
         hidden_states = self.norm_out(hidden_states, temb)
         hidden_states = self.proj_out(hidden_states)