Update code block

README.md

@@ -134,55 +134,32 @@ pip install git+github.com/huggingface/diffusers.git
 ```py
 import torch
 from collections import deque
-from diffusers import ModularPipelineBlocks, FlowMatchEulerDiscreteScheduler
 from diffusers.utils import export_to_video
+from diffusers import ModularPipelineBlocks
 from diffusers.modular_pipelines import PipelineState, WanModularPipeline
 
-class WanRTStreamingPipeline(WanModularPipeline):
-    @property
-    def default_sample_height(self):
-        return 60
-
-    @property
-    def default_sample_width(self):
-        return 104
-
-    @property
-    def frame_seq_length(self):
-        return 1560
-
-    @property
-    def seq_length(self):
-        return 32760
-
-    @property
-    def kv_cache_num_frames(self):
-        return 3
-
-    @property
-    def frame_cache_len(self):
-        return 1 + (self.kv_cache_num_frames - 1) * 4
-
-
-block_path = "krea/krea-realtime-video"
-blocks = ModularPipelineBlocks.from_pretrained(block_path, trust_remote_code=True)
-pipe = WanRTStreamingPipeline(blocks, block_path)
+repo_id = "krea/krea-realtime-video"
+blocks = ModularPipelineBlocks.from_pretrained(repo_id, trust_remote_code=True)
+pipe = WanModularPipeline(blocks, repo_id)
 
 pipe.load_components(
     trust_remote_code=True,
     device_map="cuda",
-    torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
+    torch_dtype={"default": torch.bfloat16, "vae": torch.float16},
 )
-pipe.scheduler = FlowMatchEulerDiscreteScheduler(shift=5.0)
-
-prompt = ["A cat sitting on a boat"]
 
 num_frames_per_block = 3
 num_blocks = 9
 
 frames = []
 state = PipelineState()
-state.set("frame_cache_context", deque(maxlen=pipe.frame_cache_len))
+state.set("frame_cache_context", deque(maxlen=pipe.config.frame_cache_len))
+
+prompt = ["a cat sitting on a boat"]
+
+for block in pipe.transformer.blocks:
+    block.self_attn.fuse_projections()
+
 for block_idx in range(num_blocks):
     state = pipe(
         state,
@@ -191,8 +168,9 @@ for block_idx in range(num_blocks):
         num_blocks=num_blocks,
         num_frames_per_block=num_frames_per_block,
         block_idx=block_idx,
+        generator=torch.Generator("cuda").manual_seed(42),
     )
     frames.extend(state.values["videos"][0])
 
-export_to_video(frames, "krt.mp4")
+export_to_video(frames, "output.mp4", fps=16)
 ```