Upload 5 files

cog.yaml

@@ -0,0 +1,25 @@
+# Configuration for Cog ⚙️
+# Reference: https://github.com/replicate/cog/blob/main/docs/yaml.md
+
+build:
+  gpu: true
+  system_packages:
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+  python_version: "3.11"
+  python_packages:
+    - "torch==2.0.1"
+    - "opencv-python==4.8.1.78"
+    - "torchvision==0.15.2"
+    - "pytorch_lightning==2.1.0"
+    - "einops==0.7.0"
+    - "imageio==2.31.6"
+    - "omegaconf==2.3.0"
+    - "transformers==4.35.0"
+    - "moviepy==1.0.3"
+    - "av==10.0.0"
+    - "decord==0.6.0"
+    - "kornia==0.7.0"
+    - "open-clip-torch==2.12.0"
+    - "xformers==0.0.21"
+predict: "predict.py:Predictor"

final-year-project-443dd-df6f48af0796.json

@@ -0,0 +1,13 @@
+{
+  "type": "service_account",
+  "project_id": "final-year-project-443dd",
+  "private_key_id": "df6f48af0796ab27ae03fb99d08afca2ac2b00ef",
+  "private_key": "-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCmOla5Gssdx196\n4OZyrsE1so4q3nc1fWjNs9PsQ/cm6lTHTENAMM4yHbr0no4b5jL2KgBFwAsIMAMI\nzmHJNc+r/3dnLcPOvnUH8PlkaZNpH/5eQueLz8is7QcqvtnImkg/v2wlXLXWKwWx\nlWyvW10UuYry5qsta3aclqxmhP1jem6QnQxKLiQUNdAPbqsbFyEA11QHzivsTAac\nGdDHF2V/yJ05dqRE+40EaFYbzTXHUBglC0SbgGL512KvpSC16qwFBbY9oy+jHQ55\n8uzVVw5OCSmMCI+UmOrMSe/sI67jHXgOK/GexrHNazh2XbZUSupPIIz1lsBXUl1D\n8L3UdiWVAgMBAAECggEAJwZnOcnaicE230hRkfcJESw8SEA2SG6K3lArnrpOGerF\nwIxc9YL/xbBJJgjbYB1pNXWi3r05WdC7xaN+PZjOipjNVYHfCHiaTST7x+EpZHLI\nayTV63L6r+5t0lFAG+Jst9qe7x6W6hLroUdtXrXaYnU089XHtkAWdqjBDMiIHIRO\nZM9fAnCK/0dShYa0oD1BrjrGCUDrYdJ9I3WJWU+LHBfTZfLXEWbKeE+6665bC7IY\nB9JqhMlbNJWqNwIrg/bB8lI1qIGBY7lEl32N4cQ/JXXpOtfZGx7EAlYiez+bbgnI\nbJN637gp95E8V4l1eSDoF4FdIiygVcghXavOz+AHQQKBgQDmD8NjgkZQ9iiD+1kM\nJUi5AY+xgwOPfR+/vQSM2XWe5Q2jKOR82327Hj3bgua9pWr5FlPRFOakHIohV6nx\nFHkU9LVFwA9tL2pbs+kditDwg8doJtU/wpUW9kYhJ1MAY6dyuRr53CT4XIscXlKX\nHlOK5NClSNY0wFdgIxrQ3vGR/QKBgQC4+Cb2/Chsuh2jt0mp5IESYk38f9E4/YA3\n/1m8aQIbEUfhT3Xihk/MyhOp5MisnACt4kBH2KnrFzB1FAXtAgJQMvP2hLZekTQs\nhYMD2MfsT+E1Fj/bquIh4rDmrAW2wal+HzFBcuqBo81xXrokZGood9TnDNwwow1f\nMus3AXNJeQKBgGaVqtNpWL9rNB+96TQQQAA24QMPX3wRGCIgP7IqmVcT3ePeLRw7\npzHTx1NlaEwyQaP2P8OgZUPScglyFJYqQd+FSntiq75NAUkIzS7eIlLNABLCFh7L\nPj2x7Q2Fgm5PAXCXd57oehfA9ErfCEbYP/pUE3FQLCvzhEKbBK8UanVlAoGBAIkk\nPEedmB9dMwKir/ROHsDRsD7JSgf2NK3QHumJ9ey5uFC+iIoGyX3uSfwKTBtmoz5J\nZR2f8AQFMoFr8iTS+4IY9TdPGKQvBr8H0qb0gO6eHz0sHPay0W0MVdsBqk7hcdi4\nKd375RFvsLAg6uR2qxsMFgelSlCpZA20hB9JbQAJAoGAEmCK/A7k4AJq0cWtad3y\n9wmUsvGFZUhqj1nYtZ2GchKWIcszM28G77AnT52vPNjSDfygQAVxQ7NSYIcwULiA\nMHL4pB8RQr6P4yXISh7dPG8dlrhefrm4KdVMZPOz0Cpry4KejYWKx/YMjqZxARDd\nZFRtycZMdS8kBvSHeyc4mH8=\n-----END PRIVATE KEY-----\n",
+  "client_email": "firebase-adminsdk-74lss@final-year-project-443dd.iam.gserviceaccount.com",
+  "client_id": "104174452867915111710",
+  "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+  "token_uri": "https://oauth2.googleapis.com/token",
+  "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+  "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/firebase-adminsdk-74lss%40final-year-project-443dd.iam.gserviceaccount.com",
+  "universe_domain": "googleapis.com"
+}

predict.py

@@ -0,0 +1,155 @@
+# Prediction interface for Cog ⚙️
+# https://github.com/replicate/cog/blob/main/docs/python.md
+
+
+import os
+import sys
+import argparse
+import random
+from omegaconf import OmegaConf
+from einops import rearrange, repeat
+import torch
+import torchvision
+from pytorch_lightning import seed_everything
+from cog import BasePredictor, Input, Path
+
+sys.path.insert(0, "scripts/evaluation")
+from funcs import (
+    batch_ddim_sampling,
+    load_model_checkpoint,
+    load_image_batch,
+    get_filelist,
+)
+from utils.utils import instantiate_from_config
+
+
+class Predictor(BasePredictor):
+    def setup(self) -> None:
+        """Load the model into memory to make running multiple predictions efficient"""
+
+        ckpt_path_base = "checkpoints/base_1024_v1/model.ckpt"
+        config_base = "configs/inference_t2v_1024_v1.0.yaml"
+        ckpt_path_i2v = "checkpoints/i2v_512_v1/model.ckpt"
+        config_i2v = "configs/inference_i2v_512_v1.0.yaml"
+
+        config_base = OmegaConf.load(config_base)
+        model_config_base = config_base.pop("model", OmegaConf.create())
+        self.model_base = instantiate_from_config(model_config_base)
+        self.model_base = self.model_base.cuda()
+        self.model_base = load_model_checkpoint(self.model_base, ckpt_path_base)
+        self.model_base.eval()
+
+        config_i2v = OmegaConf.load(config_i2v)
+        model_config_i2v = config_i2v.pop("model", OmegaConf.create())
+        self.model_i2v = instantiate_from_config(model_config_i2v)
+        self.model_i2v = self.model_i2v.cuda()
+        self.model_i2v = load_model_checkpoint(self.model_i2v, ckpt_path_i2v)
+        self.model_i2v.eval()
+
+    def predict(
+        self,
+        task: str = Input(
+            description="Choose the task.",
+            choices=["text2video", "image2video"],
+            default="text2video",
+        ),
+        prompt: str = Input(
+            description="Prompt for video generation.",
+            default="A tiger walks in the forest, photorealistic, 4k, high definition.",
+        ),
+        image: Path = Input(
+            description="Input image for image2video task.", default=None
+        ),
+        ddim_steps: int = Input(description="Number of denoising steps.", default=50),
+        unconditional_guidance_scale: float = Input(
+            description="Classifier-free guidance scale.", default=12.0
+        ),
+        seed: int = Input(
+            description="Random seed. Leave blank to randomize the seed", default=None
+        ),
+        save_fps: int = Input(
+            description="Frame per second for the generated video.", default=10
+        ),
+    ) -> Path:
+
+        width = 1024 if task == "text2video" else 512
+        height = 576 if task == "text2video" else 320
+        model = self.model_base if task == "text2video" else self.model_i2v
+
+        if task == "image2video":
+            assert image is not None, "Please provide image for image2video generation."
+
+        if seed is None:
+            seed = int.from_bytes(os.urandom(2), "big")
+        print(f"Using seed: {seed}")
+        seed_everything(seed)
+
+        args = argparse.Namespace(
+            mode="base" if task == "text2video" else "i2v",
+            savefps=save_fps,
+            n_samples=1,
+            ddim_steps=ddim_steps,
+            ddim_eta=1.0,
+            bs=1,
+            height=height,
+            width=width,
+            frames=-1,
+            fps=28 if task == "text2video" else 8,
+            unconditional_guidance_scale=unconditional_guidance_scale,
+            unconditional_guidance_scale_temporal=None,
+        )
+
+        ## latent noise shape
+        h, w = args.height // 8, args.width // 8
+        frames = model.temporal_length if args.frames < 0 else args.frames
+        channels = model.channels
+
+        batch_size = 1
+        noise_shape = [batch_size, channels, frames, h, w]
+        fps = torch.tensor([args.fps] * batch_size).to(model.device).long()
+        prompts = [prompt]
+        text_emb = model.get_learned_conditioning(prompts)
+
+        if args.mode == "base":
+            cond = {"c_crossattn": [text_emb], "fps": fps}
+        elif args.mode == "i2v":
+            cond_images = load_image_batch([str(image)], (args.height, args.width))
+            cond_images = cond_images.to(model.device)
+            img_emb = model.get_image_embeds(cond_images)
+            imtext_cond = torch.cat([text_emb, img_emb], dim=1)
+            cond = {"c_crossattn": [imtext_cond], "fps": fps}
+        else:
+            raise NotImplementedError
+
+        ## inference
+        batch_samples = batch_ddim_sampling(
+            model,
+            cond,
+            noise_shape,
+            args.n_samples,
+            args.ddim_steps,
+            args.ddim_eta,
+            args.unconditional_guidance_scale,
+        )
+
+        out_path = "/tmp/output.mp4"
+        vid_tensor = batch_samples[0]
+        video = vid_tensor.detach().cpu()
+        video = torch.clamp(video.float(), -1.0, 1.0)
+        video = video.permute(2, 0, 1, 3, 4)  # t,n,c,h,w
+
+        frame_grids = [
+            torchvision.utils.make_grid(framesheet, nrow=int(args.n_samples))
+            for framesheet in video
+        ]  # [3, 1*h, n*w]
+        grid = torch.stack(frame_grids, dim=0)  # stack in temporal dim [t, 3, n*h, w]
+        grid = (grid + 1.0) / 2.0
+        grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1)
+        torchvision.io.write_video(
+            out_path,
+            grid,
+            fps=args.savefps,
+            video_codec="h264",
+            options={"crf": "10"},
+        )
+        return Path(out_path)

requirements.txt

@@ -0,0 +1,24 @@
+decord==0.6.0
+einops==0.3.0
+imageio==2.9.0
+numpy==1.24.2
+omegaconf==2.1.1
+opencv_python>=4.1.2 
+pandas==2.0.0
+Pillow==9.5.0
+pytorch_lightning==1.8.3
+PyYAML==6.0
+setuptools==65.6.3
+torch==2.0.0
+torchvision>=0.7.0
+tqdm==4.65.0
+transformers==4.25.1
+moviepy>=1.0.3 
+av
+xformers
+gradio
+timm
+scikit-learn 
+open_clip_torch==2.22.0
+kornia
+sk-video>=1.1.10