From 9b46d07906325424bbe5b6a5597d0eceae3421a7 Mon Sep 17 00:00:00 2001 From: Dan Nelson Date: Fri, 1 Mar 2024 23:28:42 +0000 Subject: [PATCH 1/8] initial ci/cd --- .github/workflows/cd.yml | 86 +++++++++++++++++++++++++++++++++ .github/workflows/ci.yml | 100 +++++++++++++++++++++++++++++++++++++++ requirements_test.txt | 0 3 files changed, 186 insertions(+) create mode 100644 .github/workflows/cd.yml create mode 100644 .github/workflows/ci.yml create mode 100644 requirements_test.txt diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml new file mode 100644 index 0000000..6ad8e5c --- /dev/null +++ b/.github/workflows/cd.yml @@ -0,0 +1,86 @@ +name: replicate-cd + +on: + push: + branches: + - main + +jobs: + push: + runs-on: ubuntu-latest + env: + BASE_MODEL: 'replicate-internal/official-sdxl-prod' + PROD_MODEL: 'replicate-internal/official-sdxl-prod' + + steps: + - name: checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: what changed + id: what-changed + run: | + FILES_CHANGED=$(git diff --name-only --diff-filter=AMR ${{ github.event.before }} ${{ github.event.after }} | xargs) + echo "FILES_CHANGED=$FILES_CHANGED" >> $GITHUB_ENV + if echo "$FILES_CHANGED" | grep -q 'cog.yaml'; then + echo "cog-push=true" >> $GITHUB_OUTPUT + else + echo "cog-push=false" >> $GITHUB_OUTPUT + fi + if ${{ contains(github.event.head_commit.message, '[cog build]') }}; then + echo "cog-push=true" >> $GITHUB_OUTPUT + fi + + # if cog.yaml changes - cog build and push. else - yolo build and push! + - name: did-it-tho + env: + COG_PUSH: ${{ steps.what-changed.outputs.cog-push }} + run: | + echo "cog push?: $COG_PUSH" + echo "changed files: $FILES_CHANGED" + + - name: setup-cog + if: steps.what-changed.outputs.cog-push == 'true' + uses: replicate/setup-cog@v1.0.3 + with: + token: ${{ secrets.REPLICATE_API_TOKEN }} + install-cuda: false + + - name: cog-build + if: steps.what-changed.outputs.cog-push == 'true' + run: | + cog build + + - name: cog-push + if: steps.what-changed.outputs.cog-push == 'true' + run: | + cog push r8.im/"$PROD_MODEL" + + - name: install-yolo + run: | + sudo curl -o /usr/local/bin/yolo -L "https://github.com/replicate/yolo/releases/latest/download/yolo_$(uname -s)_$(uname -m)" + sudo chmod +x /usr/local/bin/yolo + + # yolo as hack for pushing environment variables + - name: yolo-push-env + if: steps.what-changed.outputs.cog-push == 'true' + env: + REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} + SENTRY_DSN: ${{secrets.SENTRY_DSN}} + run: | + touch meaningless_file.txt + echo "adding environment variables to $PROD_MODEL" + yolo push -e SENTRY_DSN="$SENTRY_DSN" --base $PROD_MODEL --dest $PROD_MODEL meaningless_file.txt + + - name: yolo-push + if: steps.what-changed.outputs.cog-push == 'false' + env: + REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} + SENTRY_DSN: ${{secrets.SENTRY_DSN}} + run: | + echo "pushing changes from $BASE_MODEL to $PROD_MODEL" + echo "changed files: $FILES_CHANGED" + yolo push -e SENTRY_DSN="$SENTRY_DSN" --base $BASE_MODEL --dest $PROD_MODEL $FILES_CHANGED + + ## TODO once it exists - use deployments API to update appropriate deployment diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..6a3ee54 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,100 @@ +name: replicate-ci + +on: + pull_request: + branches: + - main + +jobs: + push: + runs-on: ubuntu-latest + env: + BASE_MODEL: 'replicate-internal/official-sdxl-prod' + STAGING_MODEL: 'replicate-internal/official-sdxl-staging' + + steps: + - name: checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: what changed + id: what-changed + run: | + FILES_CHANGED=$(git diff --name-only --diff-filter=AMR ${{ github.event.before }} ${{ github.event.after }} | xargs) + echo "FILES_CHANGED=$FILES_CHANGED" >> $GITHUB_ENV + if echo "$FILES_CHANGED" | grep -q 'cog.yaml'; then + echo "cog-push=true" >> $GITHUB_OUTPUT + else + echo "cog-push=false" >> $GITHUB_OUTPUT + fi + if ${{ contains(github.event.head_commit.message, '[cog build]') }}; then + echo "cog-push=true" >> $GITHUB_OUTPUT + fi + + # if cog.yaml changes - cog build and push. else - yolo build and push! + - name: did-it-tho + env: + COG_PUSH: ${{ steps.what-changed.outputs.cog-push }} + run: | + echo "cog push?: $COG_PUSH" + echo "changed files: $FILES_CHANGED" + + - name: setup-cog + if: steps.what-changed.outputs.cog-push == 'true' + uses: replicate/setup-cog@v1.0.3 + with: + token: ${{ secrets.REPLICATE_API_TOKEN }} + install-cuda: false + + - name: cog-build + if: steps.what-changed.outputs.cog-push == 'true' + run: | + cog build + + - name: cog-push + if: steps.what-changed.outputs.cog-push == 'true' + run: | + cog push r8.im/"$STAGING_MODEL" + + - name: install-yolo + run: | + sudo curl -o /usr/local/bin/yolo -L "https://github.com/replicate/yolo/releases/latest/download/yolo_$(uname -s)_$(uname -m)" + sudo chmod +x /usr/local/bin/yolo + + # TODO: once you confirm sentry works, remove from staging. + # yolo as hack for pushing environment variables + - name: yolo-push-env + if: steps.what-changed.outputs.cog-push == 'true' + env: + REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} + SENTRY_DSN: ${{secrets.SENTRY_DSN}} + run: | + touch meaningless_file.txt + echo "adding environment variables to $STAGING_MODEL" + yolo push -e SENTRY_DSN="$SENTRY_DSN" --base $STAGING_MODEL --dest $STAGING_MODEL meaningless_file.txt + + - name: yolo-push + if: steps.what-changed.outputs.cog-push == 'false' + env: + REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} + SENTRY_DSN: ${{secrets.SENTRY_DSN}} + run: | + echo "pushing changes from $BASE_MODEL to $STAGING_MODEL" + echo "changed files: $FILES_CHANGED" + yolo push -e SENTRY_DSN="$SENTRY_DSN" --base $BASE_MODEL --dest $STAGING_MODEL $FILES_CHANGED + + - uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: install python test deps + run: | + pip install -r requirements_test.txt + + - name: test model + env: + REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} + TEST_ENV: 'staging' + run: | + pytest test_predict.py diff --git a/requirements_test.txt b/requirements_test.txt new file mode 100644 index 0000000..e69de29 From 4dd64d4e7d0c183fe989a63cae376d1d19a8c2e2 Mon Sep 17 00:00:00 2001 From: Dan Nelson Date: Fri, 1 Mar 2024 23:43:05 +0000 Subject: [PATCH 2/8] sentry added --- cog.yaml | 1 + predict.py | 346 ++++++++++++++++++++++++++++------------------------- 2 files changed, 185 insertions(+), 162 deletions(-) diff --git a/cog.yaml b/cog.yaml index 1274376..b1c3e25 100644 --- a/cog.yaml +++ b/cog.yaml @@ -24,6 +24,7 @@ build: - "fire==0.5.0" - "opencv-python>=4.1.0.25" - "mediapipe==0.10.2" + - "sentry_sdk==1.40" run: - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.0.3/pget" && chmod +x /usr/local/bin/pget diff --git a/predict.py b/predict.py index 291807a..6117e85 100644 --- a/predict.py +++ b/predict.py @@ -29,6 +29,7 @@ from safetensors import safe_open from safetensors.torch import load_file from transformers import CLIPImageProcessor +import sentry_sdk from dataset_and_utils import TokenEmbeddingsHandler @@ -158,83 +159,92 @@ def load_trained_weights(self, weights, pipe): def setup(self, weights: Optional[Path] = None): """Load the model into memory to make running multiple predictions efficient""" - start = time.time() - self.tuned_model = False - self.tuned_weights = None - if str(weights) == "weights": - weights = None - - self.weights_cache = WeightsDownloadCache() - - print("Loading safety checker...") - if not os.path.exists(SAFETY_CACHE): - download_weights(SAFETY_URL, SAFETY_CACHE) - self.safety_checker = StableDiffusionSafetyChecker.from_pretrained( - SAFETY_CACHE, torch_dtype=torch.float16 - ).to("cuda") - self.feature_extractor = CLIPImageProcessor.from_pretrained(FEATURE_EXTRACTOR) - - if not os.path.exists(SDXL_MODEL_CACHE): - download_weights(SDXL_URL, SDXL_MODEL_CACHE) - - print("Loading sdxl txt2img pipeline...") - self.txt2img_pipe = DiffusionPipeline.from_pretrained( - SDXL_MODEL_CACHE, - torch_dtype=torch.float16, - use_safetensors=True, - variant="fp16", - ) - self.is_lora = False - if weights or os.path.exists("./trained-model"): - self.load_trained_weights(weights, self.txt2img_pipe) - - self.txt2img_pipe.to("cuda") - - print("Loading SDXL img2img pipeline...") - self.img2img_pipe = StableDiffusionXLImg2ImgPipeline( - vae=self.txt2img_pipe.vae, - text_encoder=self.txt2img_pipe.text_encoder, - text_encoder_2=self.txt2img_pipe.text_encoder_2, - tokenizer=self.txt2img_pipe.tokenizer, - tokenizer_2=self.txt2img_pipe.tokenizer_2, - unet=self.txt2img_pipe.unet, - scheduler=self.txt2img_pipe.scheduler, - ) - self.img2img_pipe.to("cuda") - - print("Loading SDXL inpaint pipeline...") - self.inpaint_pipe = StableDiffusionXLInpaintPipeline( - vae=self.txt2img_pipe.vae, - text_encoder=self.txt2img_pipe.text_encoder, - text_encoder_2=self.txt2img_pipe.text_encoder_2, - tokenizer=self.txt2img_pipe.tokenizer, - tokenizer_2=self.txt2img_pipe.tokenizer_2, - unet=self.txt2img_pipe.unet, - scheduler=self.txt2img_pipe.scheduler, - ) - self.inpaint_pipe.to("cuda") - - print("Loading SDXL refiner pipeline...") - # FIXME(ja): should the vae/text_encoder_2 be loaded from SDXL always? - # - in the case of fine-tuned SDXL should we still? - # FIXME(ja): if the answer to above is use VAE/Text_Encoder_2 from fine-tune - # what does this imply about lora + refiner? does the refiner need to know about - - if not os.path.exists(REFINER_MODEL_CACHE): - download_weights(REFINER_URL, REFINER_MODEL_CACHE) - - print("Loading refiner pipeline...") - self.refiner = DiffusionPipeline.from_pretrained( - REFINER_MODEL_CACHE, - text_encoder_2=self.txt2img_pipe.text_encoder_2, - vae=self.txt2img_pipe.vae, - torch_dtype=torch.float16, - use_safetensors=True, - variant="fp16", - ) - self.refiner.to("cuda") - print("setup took: ", time.time() - start) - # self.txt2img_pipe.__class__.encode_prompt = new_encode_prompt + sentry_dsn = os.getenv("SENTRY_DSN", "no-op") + if sentry_dsn != "no-op": + sentry_sdk.init( + dsn=sentry_dsn, + ) + print("sentry init") + try: + start = time.time() + self.tuned_model = False + self.tuned_weights = None + if str(weights) == "weights": + weights = None + + self.weights_cache = WeightsDownloadCache() + + print("Loading safety checker...") + if not os.path.exists(SAFETY_CACHE): + download_weights(SAFETY_URL, SAFETY_CACHE) + self.safety_checker = StableDiffusionSafetyChecker.from_pretrained( + SAFETY_CACHE, torch_dtype=torch.float16 + ).to("cuda") + self.feature_extractor = CLIPImageProcessor.from_pretrained(FEATURE_EXTRACTOR) + + if not os.path.exists(SDXL_MODEL_CACHE): + download_weights(SDXL_URL, SDXL_MODEL_CACHE) + + print("Loading sdxl txt2img pipeline...") + self.txt2img_pipe = DiffusionPipeline.from_pretrained( + SDXL_MODEL_CACHE, + torch_dtype=torch.float16, + use_safetensors=True, + variant="fp16", + ) + self.is_lora = False + if weights or os.path.exists("./trained-model"): + self.load_trained_weights(weights, self.txt2img_pipe) + + self.txt2img_pipe.to("cuda") + + print("Loading SDXL img2img pipeline...") + self.img2img_pipe = StableDiffusionXLImg2ImgPipeline( + vae=self.txt2img_pipe.vae, + text_encoder=self.txt2img_pipe.text_encoder, + text_encoder_2=self.txt2img_pipe.text_encoder_2, + tokenizer=self.txt2img_pipe.tokenizer, + tokenizer_2=self.txt2img_pipe.tokenizer_2, + unet=self.txt2img_pipe.unet, + scheduler=self.txt2img_pipe.scheduler, + ) + self.img2img_pipe.to("cuda") + + print("Loading SDXL inpaint pipeline...") + self.inpaint_pipe = StableDiffusionXLInpaintPipeline( + vae=self.txt2img_pipe.vae, + text_encoder=self.txt2img_pipe.text_encoder, + text_encoder_2=self.txt2img_pipe.text_encoder_2, + tokenizer=self.txt2img_pipe.tokenizer, + tokenizer_2=self.txt2img_pipe.tokenizer_2, + unet=self.txt2img_pipe.unet, + scheduler=self.txt2img_pipe.scheduler, + ) + self.inpaint_pipe.to("cuda") + + print("Loading SDXL refiner pipeline...") + # FIXME(ja): should the vae/text_encoder_2 be loaded from SDXL always? + # - in the case of fine-tuned SDXL should we still? + # FIXME(ja): if the answer to above is use VAE/Text_Encoder_2 from fine-tune + # what does this imply about lora + refiner? does the refiner need to know about + + if not os.path.exists(REFINER_MODEL_CACHE): + download_weights(REFINER_URL, REFINER_MODEL_CACHE) + + print("Loading refiner pipeline...") + self.refiner = DiffusionPipeline.from_pretrained( + REFINER_MODEL_CACHE, + text_encoder_2=self.txt2img_pipe.text_encoder_2, + vae=self.txt2img_pipe.vae, + torch_dtype=torch.float16, + use_safetensors=True, + variant="fp16", + ) + self.refiner.to("cuda") + print("setup took: ", time.time() - start) + except Exception as e: + sentry_sdk.capture_exception(e) + raise e def load_image(self, path): shutil.copyfile(path, "/tmp/image.png") @@ -339,98 +349,110 @@ def predict( ), ) -> List[Path]: """Run a single prediction on the model.""" - if seed is None: - seed = int.from_bytes(os.urandom(2), "big") - print(f"Using seed: {seed}") - - if replicate_weights: - self.load_trained_weights(replicate_weights, self.txt2img_pipe) - - # OOMs can leave vae in bad state - if self.txt2img_pipe.vae.dtype == torch.float32: - self.txt2img_pipe.vae.to(dtype=torch.float16) - - sdxl_kwargs = {} - if self.tuned_model: - # consistency with fine-tuning API - for k, v in self.token_map.items(): - prompt = prompt.replace(k, v) - print(f"Prompt: {prompt}") - if image and mask: - print("inpainting mode") - sdxl_kwargs["image"] = self.load_image(image) - sdxl_kwargs["mask_image"] = self.load_image(mask) - sdxl_kwargs["strength"] = prompt_strength - sdxl_kwargs["width"] = width - sdxl_kwargs["height"] = height - pipe = self.inpaint_pipe - elif image: - print("img2img mode") - sdxl_kwargs["image"] = self.load_image(image) - sdxl_kwargs["strength"] = prompt_strength - pipe = self.img2img_pipe - else: - print("txt2img mode") - sdxl_kwargs["width"] = width - sdxl_kwargs["height"] = height - pipe = self.txt2img_pipe - - if refine == "expert_ensemble_refiner": - sdxl_kwargs["output_type"] = "latent" - sdxl_kwargs["denoising_end"] = high_noise_frac - elif refine == "base_image_refiner": - sdxl_kwargs["output_type"] = "latent" - - if not apply_watermark: - # toggles watermark for this prediction - watermark_cache = pipe.watermark - pipe.watermark = None - self.refiner.watermark = None - - pipe.scheduler = SCHEDULERS[scheduler].from_config(pipe.scheduler.config) - generator = torch.Generator("cuda").manual_seed(seed) - - common_args = { - "prompt": [prompt] * num_outputs, - "negative_prompt": [negative_prompt] * num_outputs, - "guidance_scale": guidance_scale, - "generator": generator, - "num_inference_steps": num_inference_steps, - } - - if self.is_lora: - sdxl_kwargs["cross_attention_kwargs"] = {"scale": lora_scale} - - output = pipe(**common_args, **sdxl_kwargs) - - if refine in ["expert_ensemble_refiner", "base_image_refiner"]: - refiner_kwargs = { - "image": output.images, - } + try: + if seed is None: + seed = int.from_bytes(os.urandom(2), "big") + print(f"Using seed: {seed}") + + if replicate_weights: + self.load_trained_weights(replicate_weights, self.txt2img_pipe) + + # OOMs can leave vae in bad state + if self.txt2img_pipe.vae.dtype == torch.float32: + self.txt2img_pipe.vae.to(dtype=torch.float16) + + sdxl_kwargs = {} + if self.tuned_model: + # consistency with fine-tuning API + for k, v in self.token_map.items(): + prompt = prompt.replace(k, v) + print(f"Prompt: {prompt}") + if image and mask: + print("inpainting mode") + sdxl_kwargs["image"] = self.load_image(image) + sdxl_kwargs["mask_image"] = self.load_image(mask) + sdxl_kwargs["strength"] = prompt_strength + sdxl_kwargs["width"] = width + sdxl_kwargs["height"] = height + pipe = self.inpaint_pipe + elif image: + print("img2img mode") + sdxl_kwargs["image"] = self.load_image(image) + sdxl_kwargs["strength"] = prompt_strength + pipe = self.img2img_pipe + else: + print("txt2img mode") + sdxl_kwargs["width"] = width + sdxl_kwargs["height"] = height + pipe = self.txt2img_pipe if refine == "expert_ensemble_refiner": - refiner_kwargs["denoising_start"] = high_noise_frac - if refine == "base_image_refiner" and refine_steps: - common_args["num_inference_steps"] = refine_steps + sdxl_kwargs["output_type"] = "latent" + sdxl_kwargs["denoising_end"] = high_noise_frac + elif refine == "base_image_refiner": + sdxl_kwargs["output_type"] = "latent" + + if not apply_watermark: + # toggles watermark for this prediction + watermark_cache = pipe.watermark + pipe.watermark = None + self.refiner.watermark = None + + pipe.scheduler = SCHEDULERS[scheduler].from_config(pipe.scheduler.config) + generator = torch.Generator("cuda").manual_seed(seed) + + common_args = { + "prompt": [prompt] * num_outputs, + "negative_prompt": [negative_prompt] * num_outputs, + "guidance_scale": guidance_scale, + "generator": generator, + "num_inference_steps": num_inference_steps, + } + + if self.is_lora: + sdxl_kwargs["cross_attention_kwargs"] = {"scale": lora_scale} + + output = pipe(**common_args, **sdxl_kwargs) + + if refine in ["expert_ensemble_refiner", "base_image_refiner"]: + refiner_kwargs = { + "image": output.images, + } - output = self.refiner(**common_args, **refiner_kwargs) + if refine == "expert_ensemble_refiner": + refiner_kwargs["denoising_start"] = high_noise_frac + if refine == "base_image_refiner" and refine_steps: + common_args["num_inference_steps"] = refine_steps - if not apply_watermark: - pipe.watermark = watermark_cache - self.refiner.watermark = watermark_cache + output = self.refiner(**common_args, **refiner_kwargs) - if not disable_safety_checker: - _, has_nsfw_content = self.run_safety_checker(output.images) + if not apply_watermark: + pipe.watermark = watermark_cache + self.refiner.watermark = watermark_cache - output_paths = [] - for i, image in enumerate(output.images): if not disable_safety_checker: - if has_nsfw_content[i]: - print(f"NSFW content detected in image {i}") - continue - output_path = f"/tmp/out-{i}.png" - image.save(output_path) - output_paths.append(Path(output_path)) + _, has_nsfw_content = self.run_safety_checker(output.images) + + output_paths = [] + for i, image in enumerate(output.images): + if not disable_safety_checker: + if has_nsfw_content[i]: + print(f"NSFW content detected in image {i}") + continue + output_path = f"/tmp/out-{i}.png" + image.save(output_path) + output_paths.append(Path(output_path)) + except Exception as e: + # scrub PII + del prompt + del negative_prompt + del image + del mask + + if sdxl_kwargs and 'image' in sdxl_kwargs: + image_shape = sdxl_kwargs['image'].size + sentry_sdk.capture_exception(e) + raise e if len(output_paths) == 0: raise Exception( From 645c7ea6aec0a8c5ba6ffd02931a37ff727fa8d0 Mon Sep 17 00:00:00 2001 From: Dan Nelson Date: Mon, 4 Mar 2024 04:48:28 +0000 Subject: [PATCH 3/8] working integration tests --- tests/test_predict.py | 319 ++++++++++++++++++++++-------------------- 1 file changed, 166 insertions(+), 153 deletions(-) diff --git a/tests/test_predict.py b/tests/test_predict.py index c45d2ad..787b307 100644 --- a/tests/test_predict.py +++ b/tests/test_predict.py @@ -1,189 +1,202 @@ import base64 import os -import pytest -import requests +import pickle import subprocess -import numpy as np -from PIL import Image -from threading import Thread, Lock +import sys +import time +from functools import partial from io import BytesIO -from test_utils import ( - get_image_name, - process_log_line, - capture_output, - wait_for_server_to_be_ready, -) - -# Constants -SERVER_URL = "http://localhost:5000/predictions" -HEALTH_CHECK_URL = "http://localhost:5000/health-check" - -IMAGE_NAME = "your_image_name" # replace with your image name -HOST_NAME = "your_host_name" # replace with your host name - - -@pytest.fixture(scope="session") -def server(): - image_name = get_image_name() - - command = [ - "docker", - "run", - # "-ti", - "-p", - "5000:5000", - "--gpus=all", - image_name, - ] - print("\n**********************STARTING SERVER**********************") - process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) +import numpy as np +import pytest +import replicate +import requests +from PIL import Image, ImageChops - print_lock = Lock() +ENV = os.getenv('TEST_ENV', 'local') +LOCAL_ENDPOINT = "http://localhost:5000/predictions" +MODEL = os.getenv('STAGING_MODEL', 'no model configured') - stdout_thread = Thread(target=capture_output, args=(process.stdout, print_lock)) - stdout_thread.start() +def local_run(model_endpoint: str, model_input: dict): + response = requests.post(model_endpoint, json={"input": model_input}) + data = response.json() - stderr_thread = Thread(target=capture_output, args=(process.stderr, print_lock)) - stderr_thread.start() + try: + # TODO: this will break if we test batching + datauri = data["output"][0] + base64_encoded_data = datauri.split(",")[1] + data = base64.b64decode(base64_encoded_data) + return Image.open(BytesIO(data)) + except Exception as e: + print("Error!") + print("input:", model_input) + print(data["logs"]) + raise e - wait_for_server_to_be_ready(HEALTH_CHECK_URL) - yield process +def replicate_run(model: str, version: str, model_input: dict): + output = replicate.run( + f"{model}:{version}", + input=model_input) + url = output[0] - process.terminate() - process.wait() + response = requests.get(url) + return Image.open(BytesIO(response.content)) -def test_health_check(server): - response = requests.get(HEALTH_CHECK_URL) - assert ( - response.status_code == 200 - ), f"Unexpected status code: {response.status_code}" +def wait_for_server_to_be_ready(url, timeout=300): + """ + Waits for the server to be ready. + Args: + - url: The health check URL to poll. + - timeout: Maximum time (in seconds) to wait for the server to be ready. + """ + start_time = time.time() + while True: + try: + response = requests.get(url) + data = response.json() -def get_image(response): - data = response.json() - datauri = data["output"][0] - base64_encoded_data = datauri.split(",")[1] - data = base64.b64decode(base64_encoded_data) - return Image.open(BytesIO(data)) + if data["status"] == "READY": + return + elif data["status"] == "SETUP_FAILED": + raise RuntimeError( + "Server initialization failed with status: SETUP_FAILED" + ) + except requests.RequestException: + pass -def write_image(response, output_fn): - if not os.path.exists("tmp/"): - os.makedirs("tmp") + if time.time() - start_time > timeout: + raise TimeoutError("Server did not become ready in the expected time.") - img = get_image(response) - img.save(output_fn) - return img + time.sleep(5) # Poll every 5 seconds -def roughly_the_same(img1, img2): +@pytest.fixture(scope="session") +def inference_func(): """ - Assert that pixel RGB values differ by less than 2 across an image - Handles watermarking variation + local inference uses http API to hit local server; staging inference uses python API b/c it's cleaner. """ - delta = np.array(img1, dtype=np.int32) - np.array(img2, dtype=np.int32) - return np.abs(np.mean(delta)) < 2 - - -def test_seeded_prediction(server): + if ENV == 'local': + return partial(local_run, LOCAL_ENDPOINT) + elif ENV == 'staging': + model = replicate.models.get(MODEL) + version = model.versions.list()[0] + return partial(replicate_run, MODEL, version.id) + else: + raise Exception(f"env should be local or staging but was {ENV}") + + +@pytest.fixture(scope="session", autouse=True) +def service(): + """ + Spins up local cog server to hit for tests if running locally, no-op otherwise + """ + if ENV == 'local': + print("building model") + # starts local server if we're running things locally + build_command = 'cog build -t test-model'.split() + subprocess.run(build_command, check=True) + container_name = 'cog-test' + try: + subprocess.check_output(['docker', 'inspect', '--format="{{.State.Running}}"', container_name]) + print(f"Container '{container_name}' is running. Stopping and removing...") + subprocess.check_call(['docker', 'stop', container_name]) + subprocess.check_call(['docker', 'rm', container_name]) + print(f"Container '{container_name}' stopped and removed.") + except subprocess.CalledProcessError: + # Container not found + print(f"Container '{container_name}' not found or not running.") + + run_command = f'docker run -d -p 5000:5000 --gpus all --name {container_name} test-model '.split() + process = subprocess.Popen(run_command, stdout=sys.stdout, stderr=sys.stderr) + + wait_for_server_to_be_ready("http://localhost:5000/health-check") + + yield + process.terminate() + process.wait() + stop_command = "docker stop cog-test".split() + subprocess.run(stop_command) + else: + yield + + +def image_equal_fuzzy(img_expected, img_actual, test_name='default', tol=5): + """ + Assert that average pixel values differ by less than tol across image + """ + img1 = np.array(img_expected, dtype=np.int32) + img2 = np.array(img_actual, dtype=np.int32) + + mean_delta = np.mean(np.abs(img1 - img2)) + imgs_equal = (mean_delta < tol) + if not imgs_equal: + # save failures for quick inspection + save_dir = f"tmp/{test_name}" + if not os.path.exists(save_dir): + os.makedirs(save_dir) + img_expected.save(os.path.join(save_dir, 'expected.png')) + img_actual.save(os.path.join(save_dir, 'actual.png')) + difference = ImageChops.difference(img_expected, img_actual) + difference.save(os.path.join(save_dir, 'delta.png')) + + return imgs_equal + + +def test_seeded_prediction(inference_func, request): """ SDXL w/seed should be deterministic. may need to adjust tolerance for optimized SDXLs """ data = { - "input": { - "prompt": "An astronaut riding a rainbow unicorn, cinematic, dramatic", - "num_inference_steps": 50, - "width": 1024, - "height": 1024, - "scheduler": "DDIM", - "refine": "expert_ensemble_refiner", - # Add other parameters here - "seed": 12103, - } + "prompt": "An astronaut riding a rainbow unicorn, cinematic, dramatic", + "num_inference_steps": 50, + "width": 1024, + "height": 1024, + "scheduler": "DDIM", + "refine": "expert_ensemble_refiner", + "seed": 12103, } - response = requests.post(SERVER_URL, json=data) - assert ( - response.status_code == 200 - ), f"Unexpected status code: {response.status_code}" - img_1 = get_image(response) - img_1.save("tests/assets/test_out.png") - img_2 = Image.open("tests/assets/out.png") - assert roughly_the_same(img_1, img_2) + actual_image = inference_func(data) + expected_image = Image.open("tests/assets/out.png") + assert image_equal_fuzzy(actual_image, expected_image, test_name=request.node.name) -def test_lora_load_unload(server): +def test_lora_load_unload(inference_func, request): """ - Tests generation with & without loras + Tests generation with & without loras. + This is checking for some gnarly state issues (can SDXL load / unload LoRAs), so predictions need to run in series. """ - data = { - "input": { - "prompt": "A photo of a dog on the beach", - "num_inference_steps": 50, - # Add other parameters here - "seed": 1234, - } + SEED = 1234 + base_data = { + "prompt": "A photo of a dog on the beach", + "num_inference_steps": 50, + # Add other parameters here + "seed": SEED, } - response = requests.post(SERVER_URL, json=data) - assert ( - response.status_code == 200 - ), f"Unexpected status code: {response.status_code}" - write_image(response, "tmp/base_output.png") - - data = { - "input": { - "prompt": "A photo of a TOK on the beach", - "num_inference_steps": 50, - # Add other parameters here - "replicate_weights": "https://storage.googleapis.com/dan-scratch-public/tmp/trained_model.tar", - "seed": 1234, - } + base_img_1 = inference_func(base_data) + + lora_a_data = { + "prompt": "A photo of a TOK on the beach", + "num_inference_steps": 50, + # Add other parameters here + "replicate_weights": "https://storage.googleapis.com/dan-scratch-public/sdxl/other_model.tar", + "seed": SEED } - response = requests.post(SERVER_URL, json=data) - assert ( - response.status_code == 200 - ), f"Unexpected status code: {response.status_code}" - img_1 = write_image(response, "tmp/lora_output.png") - response = requests.post(SERVER_URL, json=data) - assert ( - response.status_code == 200 - ), f"Unexpected status code: {response.status_code}" - img_2 = write_image(response, "tmp/lora_output_again.png") - - assert roughly_the_same(np.array(img_1), np.array(img_2)) + lora_a_img_1 = inference_func(lora_a_data) + assert not image_equal_fuzzy(lora_a_img_1, base_img_1, test_name=request.node.name) - data = { - "input": { - "prompt": "A photo of a TOK on the beach", - "num_inference_steps": 50, - # Add other parameters here - "replicate_weights": "https://storage.googleapis.com/dan-scratch-public/tmp/monstertoy_model.tar", - "seed": 1234, - } - } - response = requests.post(SERVER_URL, json=data) - assert ( - response.status_code == 200 - ), f"Unexpected status code: {response.status_code}" - lora_b = write_image(response, "tmp/lora_output_b.png") - assert not roughly_the_same(img_1, lora_b) + lora_a_img_2 = inference_func(lora_a_data) + assert image_equal_fuzzy(lora_a_img_1, lora_a_img_2, test_name=request.node.name) - data = { - "input": { - "prompt": "A photo of a dog on the beach", - "num_inference_steps": 50, - # Add other parameters here - "seed": 1234, - } + lora_b_data = { + "prompt": "A photo of a TOK on the beach", + "num_inference_steps": 50, + "replicate_weights": "https://storage.googleapis.com/dan-scratch-public/sdxl/monstertoy_model.tar", + "seed": SEED, } - response = requests.post(SERVER_URL, json=data) - assert ( - response.status_code == 200 - ), f"Unexpected status code: {response.status_code}" - write_image(response, "tmp/base_output_again.png") - - -if __name__ == "__main__": - pytest.main() + lora_b_img = inference_func(lora_b_data) + assert not image_equal_fuzzy(lora_a_img_1, lora_b_img, test_name=request.node.name) + assert not image_equal_fuzzy(base_img_1, lora_b_img, test_name=request.node.name) From daa2f288de0c6440aa70a1bdc8697b544517b42c Mon Sep 17 00:00:00 2001 From: Dan Nelson Date: Mon, 4 Mar 2024 22:21:37 +0000 Subject: [PATCH 4/8] small modifications to tests --- .github/workflows/ci.yml | 2 +- predict.py | 2 +- requirements_test.txt | 5 +++++ tests/test_predict.py | 5 ++++- 4 files changed, 11 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6a3ee54..61175b6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -97,4 +97,4 @@ jobs: REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} TEST_ENV: 'staging' run: | - pytest test_predict.py + pytest tests/test_predict.py diff --git a/predict.py b/predict.py index 6117e85..a60c73e 100644 --- a/predict.py +++ b/predict.py @@ -449,7 +449,7 @@ def predict( del image del mask - if sdxl_kwargs and 'image' in sdxl_kwargs: + if sdxl_kwargs is not None and 'image' in sdxl_kwargs: image_shape = sdxl_kwargs['image'].size sentry_sdk.capture_exception(e) raise e diff --git a/requirements_test.txt b/requirements_test.txt index e69de29..3184d45 100644 --- a/requirements_test.txt +++ b/requirements_test.txt @@ -0,0 +1,5 @@ +numpy +pytest +replicate +requests +Pillow \ No newline at end of file diff --git a/tests/test_predict.py b/tests/test_predict.py index 787b307..aa73709 100644 --- a/tests/test_predict.py +++ b/tests/test_predict.py @@ -83,6 +83,7 @@ def inference_func(): return partial(local_run, LOCAL_ENDPOINT) elif ENV == 'staging': model = replicate.models.get(MODEL) + print(f"model,", model) version = model.versions.list()[0] return partial(replicate_run, MODEL, version.id) else: @@ -124,9 +125,11 @@ def service(): yield -def image_equal_fuzzy(img_expected, img_actual, test_name='default', tol=5): +def image_equal_fuzzy(img_expected, img_actual, test_name='default', tol=20): """ Assert that average pixel values differ by less than tol across image + Tol determined empirically - holding everything else equal but varying seed + generates images that vary by at least 50 """ img1 = np.array(img_expected, dtype=np.int32) img2 = np.array(img_actual, dtype=np.int32) From 11bf031bff772bf040bff2794c33377e485c09ef Mon Sep 17 00:00:00 2001 From: Dan Nelson Date: Wed, 13 Mar 2024 22:30:10 +0000 Subject: [PATCH 5/8] now with updating deployments --- .github/workflows/cd.yml | 16 +++++++++++++++- .github/workflows/ci.yml | 17 +++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 6ad8e5c..8e0c71b 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -11,6 +11,7 @@ jobs: env: BASE_MODEL: 'replicate-internal/official-sdxl-prod' PROD_MODEL: 'replicate-internal/official-sdxl-prod' + PROD_DEPLOYMENT: 'replicate-internal/sdxl-prod-official-deployment' steps: - name: checkout @@ -83,4 +84,17 @@ jobs: echo "changed files: $FILES_CHANGED" yolo push -e SENTRY_DSN="$SENTRY_DSN" --base $BASE_MODEL --dest $PROD_MODEL $FILES_CHANGED - ## TODO once it exists - use deployments API to update appropriate deployment + - name: update deployment + env: + REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} + + run: + apt-get install jq + export LATEST_VERSION=$(curl -s -H "Authorization: Token $REPLICATE_API_TOKEN" \ + https://api.replicate.com/v1/models/$PROD_MODEL/versions | jq -r '.results[0].id') + curl -s \ + -X PATCH \ + -H "Authorization: Token $REPLICATE_API_TOKEN" \ + -H "Content-Type: application/json" \ + -d '{"version": "'$LATEST_VERSION'" }' \ + https://api.replicate.com/v1/deployments/$PROD_DEPLOYMENT diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 61175b6..e7b776c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -11,6 +11,7 @@ jobs: env: BASE_MODEL: 'replicate-internal/official-sdxl-prod' STAGING_MODEL: 'replicate-internal/official-sdxl-staging' + STAGING_DEPLOYMENT: 'replicate-internal/sdxl-staging-official' steps: - name: checkout @@ -98,3 +99,19 @@ jobs: TEST_ENV: 'staging' run: | pytest tests/test_predict.py + + - name: update deployment + env: + REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} + + run: + apt-get install jq + export LATEST_VERSION=$(curl -s -H "Authorization: Token $REPLICATE_API_TOKEN" \ + https://api.replicate.com/v1/models/$STAGING_MODEL/versions | jq -r '.results[0].id') + curl -s \ + -X PATCH \ + -H "Authorization: Token $REPLICATE_API_TOKEN" \ + -H "Content-Type: application/json" \ + -d '{"version": "'$LATEST_VERSION'" }' \ + https://api.replicate.com/v1/deployments/$STAGING_DEPLOYMENT + \ No newline at end of file From be09b14ff7ae827604c93cf99c312065f25c977b Mon Sep 17 00:00:00 2001 From: Dan Nelson Date: Wed, 13 Mar 2024 22:33:20 +0000 Subject: [PATCH 6/8] yamlgramming --- .github/workflows/cd.yml | 4 ++-- .github/workflows/ci.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 8e0c71b..1f4fc98 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -88,10 +88,10 @@ jobs: env: REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} - run: + run: | apt-get install jq export LATEST_VERSION=$(curl -s -H "Authorization: Token $REPLICATE_API_TOKEN" \ - https://api.replicate.com/v1/models/$PROD_MODEL/versions | jq -r '.results[0].id') + https://api.replicate.com/v1/models/$PROD_MODEL/versions | jq -r '.results[0].id') curl -s \ -X PATCH \ -H "Authorization: Token $REPLICATE_API_TOKEN" \ diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e7b776c..f963767 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -104,10 +104,10 @@ jobs: env: REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} - run: + run: | apt-get install jq export LATEST_VERSION=$(curl -s -H "Authorization: Token $REPLICATE_API_TOKEN" \ - https://api.replicate.com/v1/models/$STAGING_MODEL/versions | jq -r '.results[0].id') + https://api.replicate.com/v1/models/$STAGING_MODEL/versions | jq -r '.results[0].id') curl -s \ -X PATCH \ -H "Authorization: Token $REPLICATE_API_TOKEN" \ From 53bc60fc2608730f1a538b933d5f0f797a1a83a7 Mon Sep 17 00:00:00 2001 From: Dan Nelson Date: Wed, 13 Mar 2024 22:45:03 +0000 Subject: [PATCH 7/8] sudo yamlgramming --- .github/workflows/cd.yml | 2 +- .github/workflows/ci.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 1f4fc98..ae09167 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -89,7 +89,7 @@ jobs: REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} run: | - apt-get install jq + sudo apt-get install jq export LATEST_VERSION=$(curl -s -H "Authorization: Token $REPLICATE_API_TOKEN" \ https://api.replicate.com/v1/models/$PROD_MODEL/versions | jq -r '.results[0].id') curl -s \ diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f963767..70c7d3a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -105,7 +105,7 @@ jobs: REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}} run: | - apt-get install jq + sudo apt-get install jq export LATEST_VERSION=$(curl -s -H "Authorization: Token $REPLICATE_API_TOKEN" \ https://api.replicate.com/v1/models/$STAGING_MODEL/versions | jq -r '.results[0].id') curl -s \ From ffb30574c822a71db7767d5ffcdb8bab9b325798 Mon Sep 17 00:00:00 2001 From: Dan Nelson Date: Tue, 21 May 2024 21:45:32 +0000 Subject: [PATCH 8/8] diff everything --- .github/workflows/cd.yml | 2 +- .github/workflows/ci.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index ae09167..da76d0d 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -22,7 +22,7 @@ jobs: - name: what changed id: what-changed run: | - FILES_CHANGED=$(git diff --name-only --diff-filter=AMR ${{ github.event.before }} ${{ github.event.after }} | xargs) + FILES_CHANGED=$(git diff --name-only --diff-filter=AMR origin/main HEAD | xargs) echo "FILES_CHANGED=$FILES_CHANGED" >> $GITHUB_ENV if echo "$FILES_CHANGED" | grep -q 'cog.yaml'; then echo "cog-push=true" >> $GITHUB_OUTPUT diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 70c7d3a..e990b73 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ jobs: - name: what changed id: what-changed run: | - FILES_CHANGED=$(git diff --name-only --diff-filter=AMR ${{ github.event.before }} ${{ github.event.after }} | xargs) + FILES_CHANGED=$(git diff --name-only --diff-filter=AMR origin/main HEAD | xargs) echo "FILES_CHANGED=$FILES_CHANGED" >> $GITHUB_ENV if echo "$FILES_CHANGED" | grep -q 'cog.yaml'; then echo "cog-push=true" >> $GITHUB_OUTPUT