image generation DRAFT

Image_Generation_Dockerfile

@@ -0,0 +1,44 @@
+FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
+LABEL name="unified-io-inference"
+
+WORKDIR /root/.conda
+WORKDIR /root
+RUN apt-get update && apt-get -y install wget nano
+RUN wget \
+    https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
+    && bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \
+    && rm -f Miniconda3-latest-Linux-x86_64.sh 
+ENV PATH=/opt/conda/bin:${PATH}
+RUN bash -c "conda update -n base -c defaults conda"
+
+RUN wget -nv https://ai2-prior-uio.s3.us-west-2.amazonaws.com/public/model-weights-bin/xl_1000k.bin \
+ -O xl.bin
+RUN \
+wget -nv https://ai2-prior-uio.s3.us-west-2.amazonaws.com/public/model-weights-bin/large_1000k.bin \
+ -O large.bin
+RUN wget -nv https://ai2-prior-uio.s3.us-west-2.amazonaws.com/public/model-weights-bin/base_1000k.bin \
+ -O base.bin
+RUN \
+wget -nv https://ai2-prior-uio.s3.us-west-2.amazonaws.com/public/model-weights-bin/small_1000k.bin \
+ -O small.bin
+RUN wget -nv https://farm2.staticflickr.com/1362/1261465554_95741e918b_z.jpg -O dbg_img.png
+
+COPY uioi.yml .
+RUN bash -c "conda env create -f uioi.yml"
+COPY requirements.txt .
+RUN bash -c ". activate uioi && pip install --upgrade pip \
+ && pip install --upgrade "jax[cuda]" \
+ -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
+ && python3 -m pip install -r requirements.txt"
+
+RUN bash -c \
+". activate uioi && pip install matplotlib notebook numpy==1.23.5 nvidia-cudnn-cu11==8.6.0.163 \
+setuptools wheel spacy webdataset && python3 -m spacy download en_core_web_sm"
+ENV PYTHONPATH=/root/uio
+COPY . .
+RUN bash -c ". activate uioi && export PYTHONPATH=/root:/root/uio && python ./uio/test/check.py"
+ENV PROMPTS_FILE=./prompts.json
+ENV OUTPUT_DIR=/output
+# image labeling using alternate prompts
+ENTRYPOINT bash -c \
+". activate uioi && python ./generate_images.py xl xl.bin $PROMPTS_FILE $OUTPUT_DIR"

README.docker.md

@@ -11,6 +11,7 @@ To build using named Dockerfile (such as for VizWiz custom Dockerfile):
 ```bash
 docker build -t unified-io-inference:vizwiz -f VizWiz.Dockerfile .
 docker build -t unified-io-inference:alt-prompts -f Alt_Prompts_Dockerfile .
+docker build -t unified-io-inference:image-generation -f Image_Generation_Dockerfile .
 ```
 
 ### Run
@@ -75,3 +76,15 @@ docker run -it --gpus "device=0" \
 ...
 # bash -c ". activate uioi && python ./joiner.py /batches/batch.json /output"
 ```
+
+Image Generation
+```
+export OUTPUT_DIR=$(pwd)/output-20230508
+export PROMPTS_FILE=prompts.json
+docker run -it --gpus "device=0" \
+ -v $(pwd):/input:ro \
+ -v ${OUTPUT_DIR}:/output \
+ -e PROMPTS_FILE=/input/${PROMPTS_FILE} \
+ -e OUTPUT_DIR=/output \
+ unified-io-inference:image-generation
+```

generate_images.py

@@ -0,0 +1,63 @@
+#!python3
+# Prompt unified-io-inference for image labels
+# Usage: (activate virtual env, e.g. ". activate uioi &&"
+# python ./generate_images.py xl sl.bin $PROMPTS_FILE $OUTPUT_DIR
+import matplotlib.pyplot as plt
+import argparse
+import io
+import json
+import os
+from PIL import Image, ImageDraw, ImageFont
+from uio import runner
+from uio.configs import CONFIGS
+from uio import utils
+import numpy as np
+import spacy
+from absl import logging
+import warnings
+from itertools import islice
+from pathlib import Path
+from datetime import datetime
+# flax kicks up a lot of future warnings at the moment, ignore them
+warnings.simplefilter(action='ignore', category=FutureWarning)
+# To see INFO messages from `ModelRunner`
+logging.set_verbosity(logging.INFO)
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model_size", choices=list(CONFIGS))
+  parser.add_argument("model_weights")
+  parser.add_argument("prompts_file")
+  parser.add_argument("output_dir")
+  args = parser.parse_args()
+  output_dir = Path(args.output_dir)
+  if not output_dir.exists():
+    print(f"Output dir not found: {args.output_dir}")
+    exit()
+
+  model = runner.ModelRunner(args.model_size, args.model_weights)
+  logging.info(f"Prompts: {args.prompts_file}\nOutput Dir: {args.output_dir}")
+# read all the alternate prompts file JSON:
+  prompts_file = open(f"{args.prompts_file}")
+  prompts_data = json.load(prompts_file)
+  logging.info(f"prompts_data: {prompts_data}")
+
+  #TODO: read through prompts.  only generate single fixed image currently
+  generated = model.image_generation("red and white stripes", num_decodes=1)
+  print(f"Keys: {generated.keys()}")
+  g_score = generated["score"]
+  print(f"SCORE: {g_score}")
+  g_image = generated["image"][0]
+  print(f"Type: {type(g_image)}")
+  print(f"Shape: {g_image.shape}")
+  print(f"dtype: {g_image.dtype}")
+
+  fig, ax = plt.subplots()
+  ax.set_xticklabels([])
+  ax.set_yticklabels([])    
+  ax.imshow(g_image)
+  plt.savefig(f"{args.output_dir}/generated_image.png")
+  exit()
+
+if __name__ == "__main__":
+  main()

prompts.json

@@ -0,0 +1,5 @@
+{
+  "prompts": [
+    "A horse in a grassy field"
+  ]
+}