import torch
from diffusers import AutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModel

def tiny_pipe():
    transformer = CogVideoXTransformer3DModel(
        num_attention_heads=4, attention_head_dim=8, in_channels=4, out_channels=4,
        time_embed_dim=2, text_embed_dim=32, num_layers=1,
        sample_width=2, sample_height=2, sample_frames=9, patch_size=2,
        temporal_compression_ratio=4, max_text_seq_length=16,
    )
    vae = AutoencoderKLCogVideoX(
        in_channels=3, out_channels=3,
        down_block_types=("CogVideoXDownBlock3D",) * 4,
        up_block_types=("CogVideoXUpBlock3D",) * 4,
        block_out_channels=(8, 8, 8, 8), latent_channels=4,
        layers_per_block=1, norm_num_groups=2, temporal_compression_ratio=4,
    )
    return CogVideoXPipeline(None, None, transformer, vae, CogVideoXDDIMScheduler())

pipe = tiny_pipe()
frames = pipe(
    prompt_embeds=torch.zeros(1, 16, 32),
    height=16, width=16, num_frames=5,
    num_inference_steps=1, guidance_scale=1,
    num_videos_per_prompt=2, output_type="latent",
).frames
print(frames.shape)
assert frames.shape[0] == 2, f"expected 2 videos, got {frames.shape[0]}"

# Text-to-video: remove the forced reset.
# num_videos_per_prompt = 1

# Conditioned pipelines should either expand conditioning latents:
def _repeat_to_effective_batch(tensor, batch_size, num_videos_per_prompt):
    if tensor.shape[0] == 1:
        return tensor.repeat_interleave(batch_size * num_videos_per_prompt, dim=0)
    if tensor.shape[0] == batch_size:
        return tensor.repeat_interleave(num_videos_per_prompt, dim=0)
    return tensor

# Or reject unsupported requests until expansion is implemented:
if num_videos_per_prompt != 1:
    raise ValueError("`num_videos_per_prompt > 1` is not currently supported by this conditioned CogVideoX pipeline.")

import torch
from diffusers import AutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXFunControlPipeline, CogVideoXTransformer3DModel

transformer = CogVideoXTransformer3DModel(
    num_attention_heads=4, attention_head_dim=8, in_channels=8, out_channels=4,
    time_embed_dim=2, text_embed_dim=32, num_layers=1,
    sample_width=2, sample_height=2, sample_frames=9, patch_size=2,
    temporal_compression_ratio=4, max_text_seq_length=16,
)
vae = AutoencoderKLCogVideoX(
    in_channels=3, out_channels=3,
    down_block_types=("CogVideoXDownBlock3D",) * 4,
    up_block_types=("CogVideoXUpBlock3D",) * 4,
    block_out_channels=(8, 8, 8, 8), latent_channels=4,
    layers_per_block=1, norm_num_groups=2, temporal_compression_ratio=4,
)
pipe = CogVideoXFunControlPipeline(None, None, transformer, vae, CogVideoXDDIMScheduler())

try:
    pipe(
        prompt_embeds=torch.zeros(1, 16, 32),
        control_video_latents=torch.zeros(1, 4, 2, 2, 2),
        height=16, width=16, num_inference_steps=1,
        guidance_scale=1, output_type="latent",
    )
except Exception as e:
    print(type(e).__name__, e)

if control_video_latents is None:
    if control_video is None:
        raise ValueError("Provide either `control_video` or `control_video_latents`.")
    control_video = self.video_processor.preprocess_video(control_video, height=height, width=width)
    control_video = control_video.to(device=device, dtype=prompt_embeds.dtype)
    _, control_video_latents = self.prepare_control_latents(None, control_video)
else:
    control_video_latents = control_video_latents.to(device=device, dtype=prompt_embeds.dtype)

control_video_latents = control_video_latents.permute(0, 2, 1, 3, 4)

import torch
from diffusers import AutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModel

transformer = CogVideoXTransformer3DModel(
    num_attention_heads=4, attention_head_dim=8, in_channels=4, out_channels=4,
    time_embed_dim=2, text_embed_dim=32, num_layers=1,
    sample_width=2, sample_height=2, sample_frames=9, patch_size=2,
    temporal_compression_ratio=4, max_text_seq_length=16,
)
vae = AutoencoderKLCogVideoX(
    in_channels=3, out_channels=3,
    down_block_types=("CogVideoXDownBlock3D",) * 4,
    up_block_types=("CogVideoXUpBlock3D",) * 4,
    block_out_channels=(8, 8, 8, 8), latent_channels=4,
    layers_per_block=1, norm_num_groups=2, temporal_compression_ratio=4,
)
pipe = CogVideoXPipeline(None, None, transformer, vae, CogVideoXDDIMScheduler()).to(dtype=torch.bfloat16)

try:
    pipe(
        prompt_embeds=torch.zeros(1, 16, 32, dtype=torch.float32),
        height=16, width=16, num_frames=5,
        num_inference_steps=1, guidance_scale=1, output_type="latent",
    )
except RuntimeError as e:
    print(e)

# After prompt embedding selection in encode_prompt:
prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
if negative_prompt_embeds is not None:
    negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)

# In prepare_latents branches:
latents = latents.to(device=device, dtype=dtype)

import torch
from diffusers import AutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModel

transformer = CogVideoXTransformer3DModel(
    num_attention_heads=4, attention_head_dim=8, in_channels=4, out_channels=4,
    time_embed_dim=2, text_embed_dim=32, num_layers=1,
    sample_width=2, sample_height=2, sample_frames=9, patch_size=2,
    temporal_compression_ratio=4, max_text_seq_length=16,
)
vae = AutoencoderKLCogVideoX(
    in_channels=3, out_channels=3,
    down_block_types=("CogVideoXDownBlock3D",) * 4,
    up_block_types=("CogVideoXUpBlock3D",) * 4,
    block_out_channels=(8, 8, 8, 8), latent_channels=4,
    layers_per_block=1, norm_num_groups=2, temporal_compression_ratio=4,
)
pipe = CogVideoXPipeline(None, None, transformer, vae, CogVideoXDDIMScheduler())

try:
    pipe(
        prompt_embeds=torch.zeros(1, 16, 32),
        height=24, width=24, num_frames=5,
        num_inference_steps=1, guidance_scale=1, output_type="latent",
    )
except Exception as e:
    print(type(e).__name__, e)

spatial_multiple = self.vae_scale_factor_spatial * self.transformer.config.patch_size
if height % spatial_multiple != 0 or width % spatial_multiple != 0:
    raise ValueError(
        f"`height` and `width` have to be divisible by {spatial_multiple} "
        f"because CogVideoX patchifies VAE latents with patch_size={self.transformer.config.patch_size}; "
        f"got {height} and {width}."
    )

from diffusers import CogVideoXTransformer3DModel

model = CogVideoXTransformer3DModel(
    num_attention_heads=2, attention_head_dim=8, in_channels=4, out_channels=4,
    time_embed_dim=2, text_embed_dim=8, num_layers=1,
    sample_width=8, sample_height=8, sample_frames=8,
    patch_size=2, temporal_compression_ratio=4, max_text_seq_length=8,
)
model.set_attention_backend("native")
processors = list(model.attn_processors.values())
print([type(p).__name__ for p in processors])
print([hasattr(p, "_attention_backend") for p in processors])
assert all(hasattr(p, "_attention_backend") for p in processors)

hidden_states = dispatch_attention_fn(
    query,
    key,
    value,
    attn_mask=attention_mask,
    dropout_p=0.0,
    is_causal=False,
    backend=self._attention_backend,
    parallel_config=self._parallel_config,
)

import math
from diffusers import CogVideoXDDIMScheduler

guidance_scale = 6
num_inference_steps = 50
scheduler = CogVideoXDDIMScheduler()
scheduler.set_timesteps(num_inference_steps)

scales = [
    1 + guidance_scale * (
        (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
    )
    for t in scheduler.timesteps
]
print(scheduler.timesteps[:5].tolist())
print([round(x, 3) for x in scales[:12]], max(scales))
assert max(scales) <= guidance_scale

progress = (num_inference_steps - i) / num_inference_steps
self._guidance_scale = 1 + (guidance_scale - 1) * ((1 - math.cos(math.pi * progress**5.0)) / 2)