HAIC related changes #5

.flake8

@@ -4,12 +4,7 @@ exclude =
     notebooks/
     data/
 
-# You probably noticed the peculiar default line length.
-# Black defaults to 88 characters per line, which happens to be 10% over 80.
-# This number was found to produce significantly shorter files than sticking
-# with 80 (the most popular), or even 79 (used by the standard library).
-# In general, 90-ish seems like the wise choice.
-max-line-length = 88
+max-line-length = 120
 max-complexity = 18
 select = B,C,E,F,W,T4,B9
 import-order-style = smarkets

Makefile

@@ -2,22 +2,31 @@ sync_landmarks_model = s3cmd get --recursive --skip-existing s3://ai.h2o.wave-im
 sync_attr_models = s3cmd get --recursive --skip-existing s3://ai.h2o.wave-image-styler/public/models/stylegan2_attributes/ ./models/stylegan2_attributes/
 sync_stgan_nada_models = s3cmd get --recursive --skip-existing s3://ai.h2o.wave-image-styler/public/models/stylegan_nada/ ./models/stylegan_nada/
 sync_gfpgan_models = s3cmd get --recursive --skip-existing s3://ai.h2o.wave-image-styler/public/models/gfpgan/ ./models/gfpgan/
-download_ffhq_model = wget -P  ./models/ https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl
+download_ffhq_model = wget -nc -P ./models/ https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl
 download_sd_model  = s3cmd get --recursive --skip-existing s3://ai.h2o.wave-image-styler/public/models/stable-diffusion-v1-4/ ./models/stable_diffusion_v1_4/
 
-.PHONY: download_models
+.PHONY: download_models, download_models_haic
 
-all: download_models ## Build app
+all: download_models_haic ## Bundle app for H2O Cloud
 	h2o bundle
 
-setup: download_models ## Install dependencies
+setup: download_models ## Install dependencies for dev
 	mkdir -p var/lib/tmp/jobs
 	mkdir -p var/lib/tmp/jobs/output
 
 	python3 -m venv .venv
 	./.venv/bin/python -m pip install --upgrade pip
 	./.venv/bin/python -m pip install -r requirements_dev.txt
 
+download_models_haic:
+	mkdir -p models
+	mkdir -p models/stylegan2_attributes
+	mkdir -p models/gfpgan/
+	$(download_ffhq_model)
+	$(sync_landmarks_model)
+	$(sync_attr_models)
+	$(sync_gfpgan_models)
+
 download_models:
 	mkdir -p models
 	mkdir -p models/stylegan2_attributes

start.py

@@ -1,11 +1,15 @@
 # This additional installation is in place to help
 # resolve cmake error related to dlib.
-import subprocess
-import shlex
-import os
 import errno
+import os
+import sys
+import shlex
+import shutil
+import subprocess
 from pathlib import Path
-import logging
+from urllib.request import urlretrieve
+
+from loguru import logger as logging
 
 
 def make_dir(path: str):
@@ -22,6 +26,10 @@ def make_dir(path: str):
 PATH_ON_CLOUD = "/resources/venv/bin/python"
 # Check for CLOUD path, if doesn't exist set it to ./venv/bin/python
 PYTHON_PATH = "./.venv/bin/python" if os.path.isdir("./.venv/bin/") else PATH_ON_CLOUD
+logging = logging.patch(lambda record: record.update(name="ImageStylingArtStudio"))
+logging.add(
+    sys.stderr, format="{time} {level} {message}", filter="my_module", level="INFO"
+)
 
 logging.info("Additional installation steps.")
 # For tracking jobs
@@ -34,6 +42,26 @@ def make_dir(path: str):
 
 logging.info("Final Stage: Additional dependencies installed.")
 
+logging.info(f"Download pending model...")
+make_dir(f"{base_path}/models/stable_diffusion_v1_4/")
+urlretrieve(
+    "https://s3.amazonaws.com/ai.h2o.wave-image-styler/public/models/stable_diffusion_v1_4.zip",
+    f"{base_path}/models/stable_diffusion_v1_4.zip",
+)
+shutil.unpack_archive(
+    f"{base_path}/models/stable_diffusion_v1_4.zip",
+    f"{base_path}/",
+)
+os.remove(f"{base_path}/models/stable_diffusion_v1_4.zip")
+
+make_dir(f"{base_path}/models/stylegan_nada/")
+urlretrieve(
+    "https://s3.amazonaws.com/ai.h2o.wave-image-styler/public/models/stylegan_nada.zip",
+    f"{base_path}/models/stylegan_nada.zip",
+)
+shutil.unpack_archive(f"{base_path}/models/stylegan_nada.zip", f"{base_path}/")
+os.remove(f"{base_path}/models/stylegan_nada.zip")
+
 # Once all dependencies are installed "Start" AutoInsights.
 logging.info("Starting Image Styler.")
 DAEMON_PATH = (