sd2-512-cuda-compile-job.yaml

apiVersion: batch/v1
kind: Job
metadata:
  name: sd21-gpu-compile
spec:
  template:
    spec:
      restartPolicy: OnFailure
      nodeSelector:
        karpenter.sh/provisioner-name: amd-nvidia
      serviceAccountName: appsimulator
      volumes:
      - name: dshm
        emptyDir:
          medium: Memory
      containers:
      - name: app
        image: 133573815653.dkr.ecr.us-west-2.amazonaws.com/stablediffusion:amd64-cuda
        imagePullPolicy: Always
        volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        command: ["/start.sh"]
        args: ["compile"]
        resources:
          limits:
            nvidia.com/gpu: 1
        env:
        - name: DEVICE
          value: "cuda"
        - name: BUCKET
          value: "sdinfer"
        - name: MODEL_ID
          value: "stabilityai/stable-diffusion-2-1-base"
        - name: COMPILER_WORKDIR_ROOT
          value: "sd2_compile_dir_512"
        - name: MODEL_FILE
          value: "stable-diffusion-2-1-base"
        - name: POD_NAME
          valueFrom:
            fieldRef:
              fieldPath: metadata.name