Merge pull request #16 from mallorbc/main

Added Docker Support And Updated README For Docker

Dockerfile

@@ -0,0 +1,26 @@
+FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+ENV TZ=Europe/London
+
+SHELL ["/bin/bash", "-c"]
+
+
+RUN apt-get update && apt-get upgrade -y
+RUN apt install git -y \
+&& apt install python3 -y \
+&& apt install python3-pip -y \
+&& apt install git -y \
+&& apt install python-is-python3 -y \
+&& apt install python3-tk -y \
+&& apt install ffmpeg libsm6 libxext6  -y \ 
+&& apt install git -y
+
+
+WORKDIR /workspace
+
+RUN pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 \
+&& pip install opencv-python
+
+
+RUN pip install SwissArmyTransformer>=0.2.9 icetk gifmaker

README.md

@@ -55,10 +55,19 @@ A 4-second clip of 32 frames is shown below.
 * LocalAttention: Make sure you have CUDA installed and compile the local attention kernel.
 
 ```shell
-git clone https://github.com/Sleepychord/Image-Local-Attention
-cd Image-Local-Attention && python setup.py install
+pip install git+https://github.com/Sleepychord/Image-Local-Attention
 ```
 
+## Docker
+Alternatively you can use Docker to handle all dependencies.
+
+1. Run ```./build_image.sh```
+2. Run ```./run_image.sh```
+3. Run ```./install_image_local_attention```
+
+Optionally, after that you can recommit the image to avoid having to install image local attention again.
+
+
 ### Download
 
 Our code will automatically download or detect the models into the path defined by environment variable `SAT_HOME`. You can also manually download [CogVideo-Stage1](https://lfs.aminer.cn/misc/cogvideo/cogvideo-stage1.zip) and [CogVideo-Stage2](https://lfs.aminer.cn/misc/cogvideo/cogvideo-stage2.zip) and place them under SAT_HOME (with folders named `cogvideo-stage1` and `cogvideo-stage2`)

build_image.sh

@@ -0,0 +1,2 @@
+#!/bin/sh
+docker build -t cog .

install_image_local_attention.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+pip install git+https://github.com/Sleepychord/Image-Local-Attention

run_image.sh

@@ -0,0 +1,2 @@
+#!/bin/sh
+docker run --rm --name cog -it --gpus all  -v "${PWD}":/workspace cog