add flux docker build

Author: wwwjn

.ci/docker/README.md

@@ -20,3 +20,14 @@ It also uses the same directory structure as PyTorch.
 # Build a specific image
 ./build.sh torchtitan-ubuntu-20.04-clang12 -t myimage:latest
 ```
+
+### Flux Docker Image
+
+For flux model development, a dedicated Docker image is available that includes the necessary dependencies (transformers, einops, sentencepiece). This image is built automatically when changes are made to the flux code.
+
+To build the flux Docker image manually:
+
+```bash
+cd .ci/docker
+docker build -f ubuntu/Dockerfile.flux -t torchtitan-flux-ubuntu-20.04-clang12 .
+```

.ci/docker/build.sh

@@ -36,5 +36,4 @@ docker build \
   --build-arg "MINICONDA_VERSION=${MINICONDA_VERSION}" \
   --shm-size=1g \
   -f "${OS}"/Dockerfile \
-  "$@" \
   .

…htitan/experiments/flux/requirements.txt .ci/docker/flux-requirements.txttorchtitan/experiments/flux/requirements.txt renamed to .ci/docker/flux-requirements.txt torchtitan/experiments/flux/requirements.txt renamed to .ci/docker/flux-requirements.txt

@@ -0,0 +1,10 @@
+FROM 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/torchtitan-ubuntu-20.04-clang12:latest
+
+# Install flux dependencies
+COPY flux-requirements.txt /opt/conda/
+RUN . /opt/conda/bin/activate py_${PYTHON_VERSION} && \
+    pip install -r /opt/conda/flux-requirements.txt && \
+    rm /opt/conda/flux-requirements.txt
+
+USER ci-user
+CMD ["bash"]

.ci/docker/ubuntu/Dockerfile.flux

@@ -0,0 +1,68 @@
+name: flux-docker-builds
+
+on:
+  workflow_dispatch:
+  pull_request:
+    paths:
+      - .ci/docker/**
+      - .github/workflows/flux-docker-builds.yml
+  push:
+    branches:
+      - main
+      - release/*
+    paths:
+      - - .ci/docker/**
+      - .github/workflows/flux-docker-builds.yml
+  schedule:
+    - cron: 1 4 * * 3  # Run at 4:01 AM UTC on Wednesdays
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+jobs:
+  flux-docker-build:
+    runs-on: [self-hosted, linux.2xlarge]
+    timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - docker-image-name: torchtitan-flux-ubuntu-20.04-clang12
+            base-image-name: torchtitan-ubuntu-20.04-clang12
+    env:
+      DOCKER_IMAGE: 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/${{ matrix.docker-image-name }}
+    steps:
+      - name: Clean workspace
+        shell: bash
+        run: |
+          echo "${GITHUB_WORKSPACE}"
+          sudo rm -rf "${GITHUB_WORKSPACE}"
+          mkdir "${GITHUB_WORKSPACE}"
+
+      - name: Setup SSH (Click me for login details)
+        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        with:
+          github-secret: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Checkout the repo
+        uses: actions/checkout@v3
+
+      - name: Setup Linux
+        uses: pytorch/test-infra/.github/actions/setup-linux@main
+
+      - name: Build flux docker image
+        id: build-docker-image
+        run: |
+          cd .ci/docker
+          docker build \
+            --no-cache \
+            --progress=plain \
+            -f ubuntu/Dockerfile.flux \
+            -t 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/${{ matrix.docker-image-name }}:latest \
+            .
+          docker push 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/${{ matrix.docker-image-name }}:latest
+
+      - name: Teardown Linux
+        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        if: always()

.github/workflows/flux-docker-builds.yml

@@ -3,6 +3,12 @@
 ## Overview
 This directory contains the implementation of the [FLUX](https://github.com/black-forest-labs/flux/tree/main) model in torchtitan. In torchtitan, we showcase the pre-training process of text-to-image part of the FLUX model.
 
+## Prerequisites
+Install the required dependencies:
+```bash
+pip install -r flux-requirements.txt
+```
+
 ## Usage
 First, download the autoencoder model from HuggingFace with your own access token:
 ```bash

torchtitan/experiments/flux/README.md

@@ -60,7 +60,7 @@ def _process_cc12m_image(
         resized_img = resized_img.convert("RGB")
 
     np_img = np.array(resized_img).transpose((2, 0, 1))
-    tensor_img = torch.tensor(np_img).float() / 255.0
+    tensor_img = torch.tensor(np_img).float() / 255.0 * 2.0 - 1.0
 
     # NOTE: The following commented code is an alternative way
     # img_transform = transforms.Compose(

torchtitan/experiments/flux/dataset/flux_dataset.py

@@ -0,0 +1 @@
+.ci/docker/flux-requirements.txt