Delete inference.py

inference.py

@@ -1,369 +0,0 @@
-import torch
-from xora.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
-from xora.models.transformers.transformer3d import Transformer3DModel
-from xora.models.transformers.symmetric_patchifier import SymmetricPatchifier
-from xora.schedulers.rf import RectifiedFlowScheduler
-from xora.pipelines.pipeline_xora_video import XoraVideoPipeline
-from pathlib import Path
-from transformers import T5EncoderModel, T5Tokenizer
-import safetensors.torch
-import json
-import argparse
-from xora.utils.conditioning_method import ConditioningMethod
-import os
-import numpy as np
-import cv2
-from PIL import Image
-import random
-
-RECOMMENDED_RESOLUTIONS = [
-    (704, 1216, 41),
-    (704, 1088, 49),
-    (640, 1056, 57),
-    (608, 992, 65),
-    (608, 896, 73),
-    (544, 896, 81),
-    (544, 832, 89),
-    (512, 800, 97),
-    (512, 768, 97),
-    (480, 800, 105),
-    (480, 736, 113),
-    (480, 704, 121),
-    (448, 704, 129),
-    (448, 672, 137),
-    (416, 640, 153),
-    (384, 672, 161),
-    (384, 640, 169),
-    (384, 608, 177),
-    (384, 576, 185),
-    (352, 608, 193),
-    (352, 576, 201),
-    (352, 544, 209),
-    (352, 512, 225),
-    (352, 512, 233),
-    (320, 544, 241),
-    (320, 512, 249),
-    (320, 512, 257),
-]
-
-
-def load_vae(vae_dir):
-    vae_ckpt_path = vae_dir / "vae_diffusion_pytorch_model.safetensors"
-    vae_config_path = vae_dir / "config.json"
-    with open(vae_config_path, "r") as f:
-        vae_config = json.load(f)
-    vae = CausalVideoAutoencoder.from_config(vae_config)
-    vae_state_dict = safetensors.torch.load_file(vae_ckpt_path)
-    vae.load_state_dict(vae_state_dict)
-    if torch.cuda.is_available():
-        vae = vae.cuda()
-    return vae.to(torch.bfloat16)
-
-
-def load_unet(unet_dir):
-    unet_ckpt_path = unet_dir / "unet_diffusion_pytorch_model.safetensors"
-    unet_config_path = unet_dir / "config.json"
-    transformer_config = Transformer3DModel.load_config(unet_config_path)
-    transformer = Transformer3DModel.from_config(transformer_config)
-    unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
-    transformer.load_state_dict(unet_state_dict, strict=True)
-    if torch.cuda.is_available():
-        transformer = transformer.cuda()
-    return transformer
-
-
-def load_scheduler(scheduler_dir):
-    scheduler_config_path = scheduler_dir / "scheduler_config.json"
-    scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
-    return RectifiedFlowScheduler.from_config(scheduler_config)
-
-
-def center_crop_and_resize(frame, target_height, target_width):
-    h, w, _ = frame.shape
-    aspect_ratio_target = target_width / target_height
-    aspect_ratio_frame = w / h
-    if aspect_ratio_frame > aspect_ratio_target:
-        new_width = int(h * aspect_ratio_target)
-        x_start = (w - new_width) // 2
-        frame_cropped = frame[:, x_start : x_start + new_width]
-    else:
-        new_height = int(w / aspect_ratio_target)
-        y_start = (h - new_height) // 2
-        frame_cropped = frame[y_start : y_start + new_height, :]
-    frame_resized = cv2.resize(frame_cropped, (target_width, target_height))
-    return frame_resized
-
-
-def load_video_to_tensor_with_resize(video_path, target_height, target_width):
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        if target_height is not None:
-            frame_resized = center_crop_and_resize(
-                frame_rgb, target_height, target_width
-            )
-        else:
-            frame_resized = frame_rgb
-        frames.append(frame_resized)
-    cap.release()
-    video_np = (np.array(frames) / 127.5) - 1.0
-    video_tensor = torch.tensor(video_np).permute(3, 0, 1, 2).float()
-    return video_tensor
-
-
-def load_image_to_tensor_with_resize(image_path, target_height=512, target_width=768):
-    image = Image.open(image_path).convert("RGB")
-    image_np = np.array(image)
-    frame_resized = center_crop_and_resize(image_np, target_height, target_width)
-    frame_tensor = torch.tensor(frame_resized).permute(2, 0, 1).float()
-    frame_tensor = (frame_tensor / 127.5) - 1.0
-    # Create 5D tensor: (batch_size=1, channels=3, num_frames=1, height, width)
-    return frame_tensor.unsqueeze(0).unsqueeze(2)
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Load models from separate directories and run the pipeline."
-    )
-
-    # Directories
-    parser.add_argument(
-        "--ckpt_dir",
-        type=str,
-        required=True,
-        help="Path to the directory containing unet, vae, and scheduler subdirectories",
-    )
-    parser.add_argument(
-        "--input_video_path",
-        type=str,
-        help="Path to the input video file (first frame used)",
-    )
-    parser.add_argument(
-        "--input_image_path", type=str, help="Path to the input image file"
-    )
-    parser.add_argument(
-        "--output_path",
-        type=str,
-        default=None,
-        help="Path to save output video, if None will save in working directory.",
-    )
-    parser.add_argument("--seed", type=int, default="171198")
-
-    # Pipeline parameters
-    parser.add_argument(
-        "--num_inference_steps", type=int, default=40, help="Number of inference steps"
-    )
-    parser.add_argument(
-        "--num_images_per_prompt",
-        type=int,
-        default=1,
-        help="Number of images per prompt",
-    )
-    parser.add_argument(
-        "--guidance_scale",
-        type=float,
-        default=3,
-        help="Guidance scale for the pipeline",
-    )
-    parser.add_argument(
-        "--height",
-        type=int,
-        default=None,
-        help="Height of the output video frames. Optional if an input image provided.",
-    )
-    parser.add_argument(
-        "--width",
-        type=int,
-        default=None,
-        help="Width of the output video frames. If None will infer from input image.",
-    )
-    parser.add_argument(
-        "--num_frames",
-        type=int,
-        default=121,
-        help="Number of frames to generate in the output video",
-    )
-    parser.add_argument(
-        "--frame_rate", type=int, default=25, help="Frame rate for the output video"
-    )
-
-    parser.add_argument(
-        "--bfloat16",
-        action="store_true",
-        help="Denoise in bfloat16",
-    )
-
-    # Prompts
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        help="Text prompt to guide generation",
-    )
-    parser.add_argument(
-        "--negative_prompt",
-        type=str,
-        default="worst quality, inconsistent motion, blurry, jittery, distorted",
-        help="Negative prompt for undesired features",
-    )
-    parser.add_argument(
-        "--custom_resolution",
-        action="store_true",
-        default=False,
-        help="Enable custom resolution (not in recommneded resolutions) if specified (default: False)",
-    )
-
-    args = parser.parse_args()
-
-    if args.input_image_path is None and args.input_video_path is None:
-        assert (
-            args.height is not None and args.width is not None
-        ), "Must enter height and width for text to image generation."
-
-    # Load media (video or image)
-    if args.input_video_path:
-        media_items = load_video_to_tensor_with_resize(
-            args.input_video_path, args.height, args.width
-        ).unsqueeze(0)
-    elif args.input_image_path:
-        media_items = load_image_to_tensor_with_resize(
-            args.input_image_path, args.height, args.width
-        )
-    else:
-        media_items = None
-
-    height = args.height if args.height else media_items.shape[-2]
-    width = args.width if args.width else media_items.shape[-1]
-    assert height % 32 == 0, f"Height ({height}) should be divisible by 32."
-    assert width % 32 == 0, f"Width ({width}) should be divisible by 32."
-    assert (
-        height,
-        width,
-        args.num_frames,
-    ) in RECOMMENDED_RESOLUTIONS or args.custom_resolution, f"The selected resolution + num frames combination is not supported, results would be suboptimal. Supported (h,w,f) are: {RECOMMENDED_RESOLUTIONS}. Use --custom_resolution to enable working with this resolution."
-
-    # Paths for the separate mode directories
-    ckpt_dir = Path(args.ckpt_dir)
-    unet_dir = ckpt_dir / "unet"
-    vae_dir = ckpt_dir / "vae"
-    scheduler_dir = ckpt_dir / "scheduler"
-
-    # Load models
-    vae = load_vae(vae_dir)
-    unet = load_unet(unet_dir)
-    scheduler = load_scheduler(scheduler_dir)
-    patchifier = SymmetricPatchifier(patch_size=1)
-    text_encoder = T5EncoderModel.from_pretrained(
-        "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder"
-    )
-    if torch.cuda.is_available():
-        text_encoder = text_encoder.to("cuda")
-    tokenizer = T5Tokenizer.from_pretrained(
-        "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer"
-    )
-
-    if args.bfloat16 and unet.dtype != torch.bfloat16:
-        unet = unet.to(torch.bfloat16)
-
-    # Use submodels for the pipeline
-    submodel_dict = {
-        "transformer": unet,
-        "patchifier": patchifier,
-        "text_encoder": text_encoder,
-        "tokenizer": tokenizer,
-        "scheduler": scheduler,
-        "vae": vae,
-    }
-
-    pipeline = XoraVideoPipeline(**submodel_dict)
-    if torch.cuda.is_available():
-        pipeline = pipeline.to("cuda")
-
-    # Prepare input for the pipeline
-    sample = {
-        "prompt": args.prompt,
-        "prompt_attention_mask": None,
-        "negative_prompt": args.negative_prompt,
-        "negative_prompt_attention_mask": None,
-        "media_items": media_items,
-    }
-
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed(args.seed)
-
-    generator = torch.Generator(
-        device="cuda" if torch.cuda.is_available() else "cpu"
-    ).manual_seed(args.seed)
-
-    images = pipeline(
-        num_inference_steps=args.num_inference_steps,
-        num_images_per_prompt=args.num_images_per_prompt,
-        guidance_scale=args.guidance_scale,
-        generator=generator,
-        output_type="pt",
-        callback_on_step_end=None,
-        height=height,
-        width=width,
-        num_frames=args.num_frames,
-        frame_rate=args.frame_rate,
-        **sample,
-        is_video=True,
-        vae_per_channel_normalize=True,
-        conditioning_method=(
-            ConditioningMethod.FIRST_FRAME
-            if media_items is not None
-            else ConditioningMethod.UNCONDITIONAL
-        ),
-        mixed_precision=not args.bfloat16,
-    ).images
-
-    # Save output video
-    def get_unique_filename(base, ext, dir=".", index_range=1000):
-        for i in range(index_range):
-            filename = os.path.join(dir, f"{base}_{i}{ext}")
-            if not os.path.exists(filename):
-                return filename
-        raise FileExistsError(
-            f"Could not find a unique filename after {index_range} attempts."
-        )
-
-    for i in range(images.shape[0]):
-        # Gathering from B, C, F, H, W to C, F, H, W and then permuting to F, H, W, C
-        video_np = images[i].permute(1, 2, 3, 0).cpu().float().numpy()
-        # Unnormalizing images to [0, 255] range
-        video_np = (video_np * 255).astype(np.uint8)
-        fps = args.frame_rate
-        height, width = video_np.shape[1:3]
-        if video_np.shape[0] == 1:
-            output_filename = (
-                args.output_path
-                if args.output_path is not None
-                else get_unique_filename(f"image_output_{i}", ".png", ".")
-            )
-            cv2.imwrite(
-                output_filename, video_np[0][..., ::-1]
-            )  # Save single frame as image
-        else:
-            output_filename = (
-                args.output_path
-                if args.output_path is not None
-                else get_unique_filename(f"video_output_{i}", ".mp4", ".")
-            )
-
-            out = cv2.VideoWriter(
-                output_filename, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)
-            )
-
-            for frame in video_np[..., ::-1]:
-                out.write(frame)
-            out.release()
-
-
-if __name__ == "__main__":
-    main()