fix: optimize dockerfile to reduce build space usage

CHANGES.md

@@ -0,0 +1,27 @@
+# Changes Made to Fix Docker Build Issues
+
+## Dockerfile Optimizations
+
+1. Implemented a multi-stage build to reduce the final image size
+   - First stage (builder) installs all dependencies and downloads models
+   - Second stage (runtime) only includes necessary runtime components
+
+2. Consolidated RUN commands to reduce the number of layers
+   - Combined apt-get install commands
+   - Combined pip install commands
+   - Combined model downloads into a single layer
+
+3. Added proper cleanup after each step
+   - Added apt-get clean after installations
+   - Removed /var/lib/apt/lists/* to free up space
+   - Removed .git directories from cloned repositories
+   - Cleaned up /tmp and /var/tmp directories
+
+4. Used --no-install-recommends flag for apt-get to reduce unnecessary packages
+
+5. Optimized the final image by:
+   - Using nvidia/cuda:12.6.2-cudnn-runtime-ubuntu22.04 instead of -devel for the final stage
+   - Only copying necessary files from the builder stage
+   - Installing only runtime dependencies in the final stage
+
+These changes significantly reduce the disk space required during the build process and result in a smaller final image, which should resolve the "no space left on device" error.

Dockerfile

@@ -1,13 +1,14 @@
-FROM nvidia/cuda:12.6.2-cudnn-devel-ubuntu22.04
+FROM nvidia/cuda:12.6.2-cudnn-devel-ubuntu22.04 AS builder
 
 WORKDIR /
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=True
 ENV TERM=xterm-256color
 
-# Install basic system dependencies and Git
-RUN apt update -y && apt install -y \
+# Install basic system dependencies and Git in a single layer with aggressive cleanup
+RUN apt-get update -y && \
+    apt-get install -y --no-install-recommends \
     software-properties-common \
     build-essential \
     libgl1 \
@@ -23,9 +24,6 @@ RUN apt update -y && apt install -y \
     python3.10 \
     python3.10-venv \
     python3-pip \
-    && add-apt-repository -y ppa:git-core/ppa \
-    && apt update -y \
-    && apt install -y \
     git \
     git-lfs \
     sudo \
@@ -35,67 +33,105 @@ RUN apt update -y && apt install -y \
     wget \
     unzip \
     unrar \
-    ffmpeg \
-    && rm -rf /var/lib/apt/lists/*
+    ffmpeg && \
+    add-apt-repository -y ppa:git-core/ppa && \
+    apt-get update -y && \
+    apt-get install -y --no-install-recommends git && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
 # Set Python 3.10 as default
 RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
     update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1
 
-# Verify Python version
-RUN python --version && pip --version
-
-# Install PyTorch and related packages
+# Install PyTorch and related packages with cleanup
 RUN pip install --no-cache-dir \
     torch==2.5.0+cu124 \
     torchvision==0.20.0+cu124 \
     torchaudio==2.5.0+cu124 \
+    --extra-index-url https://download.pytorch.org/whl/cu124 && \
+    pip install --no-cache-dir \
     torchtext==0.18.0 \
     torchdata==0.8.0 \
-    --extra-index-url https://download.pytorch.org/whl/cu124
-
-# Install AI/ML specific packages
-RUN pip install --no-cache-dir \
     xformers==0.0.28.post2 \
     https://github.com/camenduru/wheels/releases/download/torch-2.5.0-cu124/flash_attn-2.6.3-cp310-cp310-linux_x86_64.whl \
     opencv-python \
     imageio \
     imageio-ffmpeg \
     ffmpeg-python \
     av \
-    runpod \
+    runpod==1.7.13 \
     torchsde \
     einops \
     diffusers \
     transformers \
     accelerate \
-    git+https://github.com/timpietrusky/vercel_blob.git
+    git+https://github.com/timpietrusky/vercel_blob.git && \
+    rm -rf /root/.cache/pip
 
-# Clone ComfyUI and custom nodes
-RUN git clone https://github.com/comfyanonymous/ComfyUI /ComfyUI && \
-    git clone -b dev https://github.com/camenduru/ComfyUI-MochiWrapper /ComfyUI/custom_nodes/ComfyUI-MochiWrapper && \
-    git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite /ComfyUI/custom_nodes/ComfyUI-VideoHelperSuite && \
-    git clone https://github.com/kijai/ComfyUI-KJNodes /ComfyUI/custom_nodes/ComfyUI-KJNodes && \
-    git clone https://github.com/ltdrdata/ComfyUI-Manager /ComfyUI/custom_nodes/ComfyUI-Manager && \
-    git clone https://github.com/ciri/comfyui-model-downloader /ComfyUI/custom_nodes/comfyui-model-downloader
+# Clone ComfyUI and custom nodes with immediate cleanup
+RUN git clone --depth 1 https://github.com/comfyanonymous/ComfyUI /ComfyUI && \
+    git clone --depth 1 -b dev https://github.com/camenduru/ComfyUI-MochiWrapper /ComfyUI/custom_nodes/ComfyUI-MochiWrapper && \
+    git clone --depth 1 https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite /ComfyUI/custom_nodes/ComfyUI-VideoHelperSuite && \
+    git clone --depth 1 https://github.com/kijai/ComfyUI-KJNodes /ComfyUI/custom_nodes/ComfyUI-KJNodes && \
+    git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Manager /ComfyUI/custom_nodes/ComfyUI-Manager && \
+    git clone --depth 1 https://github.com/ciri/comfyui-model-downloader /ComfyUI/custom_nodes/comfyui-model-downloader && \
+    find /ComfyUI -name ".git" -type d -exec rm -rf {} + 2>/dev/null || true
 
-# Download model files
+# Create model directories
 RUN mkdir -p /ComfyUI/models/diffusion_models/mochi && \
     mkdir -p /ComfyUI/models/vae/mochi && \
-    mkdir -p /ComfyUI/models/clip && \
-    aria2c --console-log-level=error -c -x 16 -s 16 -k 1M \
+    mkdir -p /ComfyUI/models/clip
+
+# Download model files one by one with cleanup after each download
+RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M \
     https://huggingface.co/Kijai/Mochi_preview_comfy/resolve/main/mochi_preview_dit_bf16.safetensors \
     -d /ComfyUI/models/diffusion_models/mochi -o mochi_preview_dit_bf16.safetensors && \
+    rm -rf /tmp/* /var/tmp/* && \
     aria2c --console-log-level=error -c -x 16 -s 16 -k 1M \
     https://huggingface.co/Kijai/Mochi_preview_comfy/resolve/main/mochi_preview_vae_decoder_bf16.safetensors \
     -d /ComfyUI/models/vae/mochi -o mochi_preview_vae_decoder_bf16.safetensors && \
+    rm -rf /tmp/* /var/tmp/* && \
     aria2c --console-log-level=error -c -x 16 -s 16 -k 1M \
     https://huggingface.co/mcmonkey/google_t5-v1_1-xxl_encoderonly/resolve/main/model.safetensors \
-    -d /ComfyUI/models/clip -o google_t5-v1_1-xxl_encoderonly-fp16.safetensors
+    -d /ComfyUI/models/clip -o google_t5-v1_1-xxl_encoderonly-fp16.safetensors && \
+    rm -rf /tmp/* /var/tmp/*
+
+# Start with a fresh image for the final stage
+FROM nvidia/cuda:12.6.2-cudnn-runtime-ubuntu22.04
+
+WORKDIR /
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=True
+ENV TERM=xterm-256color
+
+# Install only the necessary runtime dependencies
+RUN apt-get update -y && \
+    apt-get install -y --no-install-recommends \
+    python3.10 \
+    python3.10-venv \
+    python3-pip \
+    libgl1 \
+    libglib2.0-0 \
+    ffmpeg \
+    curl \
+    wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Set Python 3.10 as default
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
+    update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1
+
+# Copy only the necessary files from the builder stage
+COPY --from=builder /ComfyUI /ComfyUI
+COPY --from=builder /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages
 
 # Copy source code and start script
 COPY ./src /ComfyUI/src
 COPY ./src/start.sh /start.sh
+COPY ./test_input.json /test_input.json
 RUN chmod +x /start.sh
 
 ENTRYPOINT ["/start.sh"]

README.md

@@ -1,12 +1,12 @@
-# worker-mochi
+# Mochi Worker for RunPod
 
-> Generate videos with Mochi as an endpoint on RunPod
+This worker provides a RunPod serverless endpoint for Mochi, a text-to-video model from Kijai.
 
 ## Features
 
-- Video generation using [Mochi 1](https://github.com/genmoai/mochi) by [Genmo](https://genmo.ai)
-- Automatic model loading and initialization
-- [UploadThing](https://uploadthing.com/) integration for video upload
+- Generate videos from text prompts using Mochi
+- Customize video dimensions, number of frames, and other parameters
+- Optimized for performance on RunPod infrastructure
 
 ## API Reference
 
@@ -36,39 +36,47 @@
 }
 ```
 
-#### Core Parameters
-
-| Parameter         | Description                                                            | Default |
-| ----------------- | ---------------------------------------------------------------------- | ------- |
-| `positive_prompt` | Text description of what you want to generate                          | `""`    |
-| `negative_prompt` | Text description of what you want to avoid in the generation           | `""`    |
-| `width`           | Output video width in pixels                                           | `848`   |
-| `height`          | Output video height in pixels                                          | `480`   |
-| `seed`            | Random seed for reproducible results                                   | `1337`  |
-| `steps`           | Number of denoising steps (higher = better quality, slower generation) | `40`    |
-| `cfg`             | Classifier-free guidance scale (how closely to follow the prompt)      | `6`     |
-| `num_frames`      | Number of frames to generate                                           | `31`    |
-
-#### VAE Parameters
-
-| Parameter                    | Description                                | Default |
-| ---------------------------- | ------------------------------------------ | ------- |
-| `enable_vae_tiling`          | Enable tiling for VAE decoding             | `false` |
-| `tile_sample_min_width`      | Minimum tile width when tiling is enabled  | `312`   |
-| `tile_sample_min_height`     | Minimum tile height when tiling is enabled | `160`   |
-| `tile_overlap_factor_width`  | Overlap factor between tiles (width)       | `0.25`  |
-| `tile_overlap_factor_height` | Overlap factor between tiles (height)      | `0.25`  |
-| `auto_tile_size`             | Automatically determine tile size          | `false` |
-| `frame_batch_size`           | Number of frames to process in parallel    | `8`     |
+### Output
+
+The worker returns a URL to the generated video.
+
+## Local Development
+
+### Prerequisites
+
+- Docker
+- NVIDIA GPU with CUDA support
+- Git LFS
+
+### Setup
+
+1. Clone this repository
+2. Copy `.env.example` to `.env` and configure as needed
+3. Run `docker-compose up --build`
 
 ## Deployment
 
-Deploy this worker on RunPod using the [GitHub Integration](https://docs.runpod.io/serverless/github-integration).
+### RunPod
+
+1. Create a new serverless template on RunPod
+2. Use the Docker image `runpod/worker-mochi:latest`
+3. Configure the template with appropriate GPU resources
+4. Deploy the endpoint
+
+## Models
 
-## Development
+This worker uses the following models:
 
-For development and contribution guidelines, please see our [Contributing Guide](.github/CONTRIBUTING.md).
+- Mochi Preview DIT (BF16)
+- Mochi Preview VAE Decoder (BF16)
+- Google T5-v1.1-XXL Encoder Only (FP16)
 
 ## License
 
-[MIT License](LICENSE)
+See the [LICENSE](LICENSE) file for details.
+
+## Acknowledgements
+
+- [Kijai](https://github.com/kijai) for creating Mochi
+- [ComfyUI](https://github.com/comfyanonymous/ComfyUI) for the backend framework
+- [RunPod](https://runpod.io) for the serverless infrastructure