inference-sh/CausVid
Updated
I see both first block cache and teacache referenced. Does this implementation also cover (can be used in place) for teacache? Thanks and amazing PR!
import torch
from diffusers import CogView4Pipeline
from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
pipe = CogView4Pipeline.from_pretrained("THUDM/CogView4-6B", torch_dtype=torch.bfloat16)
pipe.to("cuda")
apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
prompt = "A photo of an astronaut riding a horse on mars"
image = pipe(prompt, generator=torch.Generator().manual_seed(42)).images[0]
image.save("output.png")