diffusers==0.30.1
peft==0.17.1
transformers==4.44
imageio
imageio-ffmpeg==0.5.1
pandas
wandb
accelerate
matplotlib
sentencepiece
yt-dlp
datasets
torch
torchvision
av
ffmpeg

import torch
import torch.nn as nn
import numpy as np
import imageio
from PIL import Image
import torchvision.transforms as T
from diffusers import CogVideoXPipeline
from huggingface_hub import PyTorchModelHubMixin
from peft import LoraConfig, get_peft_model, TaskType
from diffusers import CogVideoXTransformer3DModel


class TruncatedCogVideoX(nn.Module):
    def __init__(self, n_blocks=6, out_dim=1, reward_backbone=False):
        super().__init__()
        self.reward_backbone = reward_backbone

        pretrained_model = CogVideoXTransformer3DModel.from_pretrained(
            "zai-org/CogVideoX-2b",
            subfolder='transformer',
            torch_dtype=torch.float32
        )
        # Use pretrained CogVideoX backbone
        self.patch_embed = pretrained_model.patch_embed
        self.embedding_dropout = pretrained_model.embedding_dropout
        self.time_proj = pretrained_model.time_proj
        self.time_embedding = pretrained_model.time_embedding
        self.config = pretrained_model.config
        self.register_buffer("pos_embedding", pretrained_model.pos_embedding, persistent=False)

        self.gradient_checkpointing = False

        # Keep first n blocks
        self.transformer_blocks = nn.ModuleList(
            [pretrained_model.transformer_blocks[i] for i in range(n_blocks)]
        )
        
        # Norm after transformer
        self.norm_final = pretrained_model.norm_final
        #self.norm_out = pretrained_model.norm_out
        #self.proj_out = pretrained_model.proj_out


        # Classification head
        self.pool = nn.AdaptiveAvgPool1d(1)   # global average pooling over sequence
        
        if reward_backbone:
            self.fc = None
        else:
            self.fc = nn.Linear(1920, out_dim)    # backbone hidden dim = 1920
        
        del pretrained_model

    def forward(self, hidden_states, timesteps, encoder_hidden_states, image_rotary_emb=None):
        """
        hidden_states: (B, C, T, H, W) video input tensor
        """
        batch_size, num_frames, channels, height, width = hidden_states.shape

        # Time embeddings
        t_emb = self.time_proj(timesteps).to(dtype=torch.float16)
        emb = self.time_embedding(t_emb, None)

        # Patch embeddings
        hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)

        # Position embedding
        text_seq_length = encoder_hidden_states.shape[1]
        
        if not self.config.use_rotary_positional_embeddings:
            seq_length = height * width * num_frames // (self.config.patch_size**2)

            pos_embeds = self.pos_embedding[:, : text_seq_length + seq_length]
            hidden_states = hidden_states + pos_embeds
            hidden_states = self.embedding_dropout(hidden_states)

        encoder_hidden_states = hidden_states[:, :text_seq_length]
        hidden_states = hidden_states[:, text_seq_length:]


        for i, block in enumerate(self.transformer_blocks):
            if self.gradient_checkpointing:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)

                    return custom_forward

                ckpt_kwargs = {"use_reentrant": False}
                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states,
                    encoder_hidden_states,
                    emb,
                    image_rotary_emb,
                    **ckpt_kwargs,
                )
            else:
                hidden_states, encoder_hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    temb=emb,
                    image_rotary_emb=image_rotary_emb,
                )
            
        if not self.config.use_rotary_positional_embeddings:
            # CogVideoX-2B
            hidden_states = self.norm_final(hidden_states)
        else:
            # CogVideoX-5B
            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
            hidden_states = self.norm_final(hidden_states)
            hidden_states = hidden_states[:, text_seq_length:]
        

        hidden_states = hidden_states.view(batch_size, num_frames, -1, 1920).permute(0, 1, 3, 2)
        hidden_states = self.pool(hidden_states.reshape(batch_size, num_frames * 1920, -1))

        if self.reward_backbone:
            return hidden_states.view(batch_size, num_frames, 1920)[:, -1, :]
        
        output = self.fc(hidden_states.view(batch_size, num_frames, 1920))

        return output


class LoRACogVideoX(TruncatedCogVideoX):
    def __init__(self, n_blocks, lora_layers, out_dim=1, r=8, lora_alpha=16, reward_backbone=False):
        super().__init__(n_blocks=n_blocks, out_dim=out_dim, reward_backbone=reward_backbone)
        target_modules = []
        for i in lora_layers:
            if i < len(self.transformer_blocks):
                target_modules.extend([
                    f"transformer_blocks.{i}.attn1.to_q", f"transformer_blocks.{i}.attn1.to_k",
                    f"transformer_blocks.{i}.attn1.to_v", f"transformer_blocks.{i}.attn1.to_out.0",
                    f"transformer_blocks.{i}.attn2.to_q", f"transformer_blocks.{i}.attn2.to_k",
                    f"transformer_blocks.{i}.attn2.to_v", f"transformer_blocks.{i}.attn2.to_out.0",
                ])
        lora_config = LoraConfig(
            r=r, lora_alpha=lora_alpha, target_modules=target_modules,
            lora_dropout=0.05, bias="none", task_type=TaskType.FEATURE_EXTRACTION
        )
        self = get_peft_model(self, lora_config)


class ModelWrapper(nn.Module, PyTorchModelHubMixin):
    def __init__(self, model_type, model_args, num_frames, mode='linear',
                 reward_hidden_size=128, reward_final_size=3, use_reward_layer=False):
        super().__init__()
        self.backbone = LoRACogVideoX(*model_args)
        self.use_reward_layer = use_reward_layer
        if use_reward_layer:
            self.reward_layer = nn.Sequential(
                nn.Linear(1920, reward_hidden_size),
                nn.SiLU(),
                nn.Linear(reward_hidden_size, reward_final_size),
            )
        self.use_linear = (mode == 'linear')
        if self.use_linear:
            self.linear = nn.Linear(reward_final_size, 1)

    def forward(self, hidden_states, timesteps, encoder_hidden_states, image_rotary_emb=None):
        out = self.backbone(hidden_states, timesteps, encoder_hidden_states, image_rotary_emb)
        if self.use_reward_layer:
            out = self.reward_layer(out.squeeze(-1))
        if self.use_linear:
            out = self.linear(out)
        return out





def video_to_latent_tensor(
    video_path,
    prompt,
    pipe=None,
    device="cuda",
    h=480,
    w=720,
    num_frames=49,
    dtype=torch.float16
):
    """
    Convert a single video into VAE latent tensor.

    Args:
        video_path (str): path to video file
        pipe (CogVideoXPipeline, optional): pass preloaded pipeline to avoid reloading
        device (str): cuda / cpu
        h, w (int): resize dimensions
        num_frames (int): number of frames to sample
        dtype: torch dtype

    Returns:
        latents (torch.Tensor): shape [1, C, T, H', W']
    """

    # Load pipeline if not provided (better to reuse externally)
    if pipe is None:
        pipe = CogVideoXPipeline.from_pretrained(
            "THUDM/CogVideoX-2b",
            torch_dtype=dtype
        ).to(device)

    vae = pipe.vae
    text_encoder = pipe.text_encoder.eval()
    tokenizer = pipe.tokenizer
    vae.eval()

    transform = T.Compose([
        T.Resize((h, w)),
        T.ToTensor(),
        T.Normalize([0.5], [0.5])
    ])

    # --- Read video ---
    reader = imageio.get_reader(video_path, format="ffmpeg")
    frames = [Image.fromarray(frame) for frame in reader]
    reader.close()

    total_frames = len(frames)
    if total_frames == 0:
        raise ValueError("Video has no frames")

    # --- Frame sampling ---
    if total_frames < num_frames:
        indices = np.linspace(0, total_frames - 1, total_frames, dtype=int)
    else:
        indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)

    sampled_frames = [transform(frames[i]) for i in indices]

    # --- Tensor shape: [1, C, T, H, W] ---
    video_tensor = torch.stack(sampled_frames, dim=1).unsqueeze(0).to(device, dtype=dtype)

    # --- Encode ---
    with torch.no_grad():
        latents = vae.encode(video_tensor).latent_dist.sample()
        latents = latents * vae.config.scaling_factor
        
        input_ids = tokenizer(
            prompt,
            padding="max_length",
            max_length=226,
            truncation=True,
            add_special_tokens=True,
            return_tensors="pt",
        ).input_ids
        prompt_embeds = text_encoder(input_ids.cuda())[0]

    return latents, prompt_embeds


if __name__ == '__main__':
    # Load the model
    video_path = 'common_videos_gen3_instances/instance_00000/gen3_00102.mp4'
    prompt = 'A ball rolling on a table'
    video_input, prompt_embeds = video_to_latent_tensor(video_path, prompt)
    
    
    torch.cuda.empty_cache()
    
    timestep = torch.zeros((video_input.shape[0],), dtype=torch.long, device=video_input.device)
    video_input = video_input.permute(0, 2, 1, 3, 4)
    
    model = ModelWrapper.from_pretrained("sasuke-ss1/GT-SVJ")
    model.half().cuda()
    model.eval()
    
    with torch.no_grad():
        output = model(video_input, timestep, prompt_embeds)
    print(output)