import torch

from diffusers import LTXPipeline, LTXImageToVideoPipeline, LTXVideoTransformer3DModel, GGUFQuantizationConfig, AutoencoderKLLTXVideo


device = "cuda" if torch.cuda.is_available() else "cpu"

ckpt_path = (
    "https://huggingface.co/city96/LTX-Video-gguf/blob/main/ltx-video-2b-v0.9-Q3_K_S.gguf"
)
transformer = LTXVideoTransformer3DModel.from_single_file(
    ckpt_path,
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.float32),
    torch_dtype=torch.float32,
)
vae = AutoencoderKLLTXVideo.from_single_file("https://huggingface.co/Lightricks/LTX-Video/ltx-video-2b-v0.9.safetensors", torch_dtype=torch.float32)
t2v_pipe = LTXPipeline.from_pretrained(
    "Lightricks/LTX-Video",
    transformer=transformer,
    vae=vae,
    torch_dtype=torch.float32,
)
t2v_pipe.to(device)

i2v_pipe = LTXImageToVideoPipeline.from_pipe(t2v_pipe)
i2v_pipe.to(device)