sync some work (next: fix pytorch conflict)

gadicc · gadicc · commit 9927fa6edbd7 · 2022-11-12T10:08:32.000Z
diff --git a/Dockerfile b/Dockerfile
@@ -50,19 +50,24 @@ FROM base as output
 RUN mkdir /api
 WORKDIR /api
 
+## XXXX playing around a lot.
+# pip installs pytorch 1.13 and uninstalls 1.12 (needed by xformers)
+# recomment conda update; didn't help.  need to solve above issue.
+
+RUN conda update -n base -c defaults conda
 # We need python 3.9 or 3.10 for xformers
 # Yes, we install pytorch twice... will switch base image in future
-# RUN conda update -n base -c defaults conda
 RUN conda create -n xformers python=3.10
 SHELL ["/opt/conda/bin/conda", "run", "--no-capture-output", "-n", "xformers", "/bin/bash", "-c"]
 RUN python --version
 RUN conda install -c pytorch -c conda-forge cudatoolkit=11.6 pytorch=1.12.1
 RUN conda install xformers -c xformers/label/dev
 
 # Install python packages
-RUN pip3 install --upgrade pip
+# RUN pip3 install --upgrade pip
+RUN https_proxy="" REQUESTS_CA_BUNDLE="" conda install pip
 ADD requirements.txt requirements.txt
-RUN pip3 install -r requirements.txt
+RUN pip install -r requirements.txt
 
 # Required to build flash attention
 # Turing: 7.5 (RTX 20s, Quadro), Ampere: 8.0 (A100), 8.6 (RTX 30s)
@@ -119,7 +124,13 @@ ARG USE_PATCHMATCH=0
 RUN if [ "$USE_PATCHMATCH" = "1" ] ; then apt-get install -yqq python3-opencv ; fi
 COPY --from=patchmatch /tmp/PyPatchMatch PyPatchMatch
 
+ARG USE_DREAMBOOTH=1
+RUN if [ "$USE_DREAMBOOTH" = "1" ] ; then \
+    pip install -r diffusers/examples/dreambooth/requirements.txt ; \
+  fi
+
 # Add your custom app code, init() and inference()
+ADD train_dreambooth.py .
 ADD send.py .
 ADD app.py .
 
diff --git a/app.py b/app.py
@@ -22,6 +22,7 @@
 import skimage.measure
 from PyPatchMatch import patch_match
 from getScheduler import getScheduler, SCHEDULERS
+from train_dreambooth import TrainDreamBooth
 import re
 
 MODEL_ID = os.environ.get("MODEL_ID")
@@ -116,6 +117,10 @@ def truncateInputs(inputs: dict):
         for item in ["init_image", "mask_image", "image"]:
             if item in modelInputs:
                 modelInputs[item] = modelInputs[item][0:6] + "..."
+        if "instance_images" in modelInputs:
+            modelInputs["instance_images"] = list(
+                map(lambda str: str[0:6] + "...", modelInputs["instance_images"])
+            )
     return clone
 
 
@@ -209,6 +214,14 @@ def inference(all_inputs: dict) -> dict:
             model_inputs.get("mask_image"), "mask_image"
         )
 
+    if "instance_images" in model_inputs:
+        model_inputs["instance_images"] = list(
+            map(
+                lambda str: decodeBase64Image(str, "instance_image"),
+                model_inputs["instance_images"],
+            )
+        )
+
     seed = model_inputs.get("seed", None)
     if seed == None:
         generator = torch.Generator(device="cuda")
@@ -259,6 +272,14 @@ def inference(all_inputs: dict) -> dict:
     # with autocast("cuda"):
     # image = pipeline(**model_inputs).images[0]
 
+    if call_inputs.get("train", None) == "dreambooth":
+        result = TrainDreamBooth(pipeline, model_inputs)
+        send("inference", "done", {"startRequestId": startRequestId})
+        inferenceTime = get_now() - inferenceStart
+        timings = {"init": initTime, "inference": inferenceTime}
+        result.update({"timings": timings})
+        return result
+
     with torch.inference_mode():
         # autocast im2img and inpaint which are broken in 0.4.0, 0.4.1
         # still broken in 0.5.1
diff --git a/test.py b/test.py
@@ -266,7 +266,7 @@ def main(tests_to_run, banana, extraCallInputs):
     name = "dreambooth"
     inputs = {
         "modelInputs": {
-            "prompt": "girl with a pearl earing standing in a big room",
+            "instance_prompt": "a photo of sks dog",
             "instance_images": list(
                 map(b64encode_file, list(Path("tests/fixtures/dreambooth").iterdir()))
             ),
@@ -276,29 +276,12 @@ def main(tests_to_run, banana, extraCallInputs):
             "MODEL_ID": "runwayml/stable-diffusion-v1-5",
             "PIPELINE": "StableDiffusionPipeline",
             "SCHEDULER": "LMSDiscreteScheduler",
+            "train": "dreambooth",
         },
     }
-    print(inputs)
 
-    # print(json.dumps(inputs, indent=4))
-    exit
-
-    """
     print("Running test: " + name)
     response = requests.post("http://localhost:8000/", json=inputs)
     result = response.json()
 
-    if result.get("images_base64", None) == "None":
-        print(json.dumps(result, indent=4))
-        print()
-        exit
-
-    images_base64 = result.get("images_base64", None)
-    if images_base64:
-        for idx, image_byte_string in enumerate(images_base64):
-            decode_and_save(image_byte_string, f"{name}_{idx}")
-    else:
-        decode_and_save(result["image_base64"], name)
-
-    print()
-    """
+    print(result)
diff --git a/train_dreambooth.py b/train_dreambooth.py
