cog-safe-push with test cases

replicate · Sep 11, 2024 · 471df8e · 471df8e
1 parent 3d03901
commit 471df8e
Show file tree

Hide file tree

Showing 2 changed files with 122 additions and 37 deletions.
diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml
@@ -4,18 +4,10 @@ on:
   workflow_dispatch:
     branches: [main]
     inputs:
-      test_only:
+      no_push:
         description: 'Test only, without pushing to prod'
         type: boolean
         default: true
-      compare_outputs:
-        description: 'Compare hotswap prediction outputs between existing version and new version'
-        type: boolean
-        default: true
-      compare_trainer_outputs:
-        description: 'Compare training outputs between existing version and new version'
-        type: boolean
-        default: true
 
 jobs:
   cog-safe-push:
@@ -46,35 +38,9 @@ jobs:
       run: |
         pip install git+https://github.com/replicate/cog-safe-push.git
 
-    - name: Run cog-safe-push to test the trainer
-      if: ${{ github.event.inputs.compare_trainer_outputs == 'true' }}
-      env:
-        ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
-        REPLICATE_API_TOKEN: ${{ secrets.REPLICATE_API_TOKEN }}
-      run: |
-        cog-safe-push -vv \
-          --test-model=replicate-internal/test-flux-fine-tuner \
-          --test-only \
-          --train \
-          --train-destination=replicate-internal/test-flux-fine-tuner-dest \
-          --predict-timeout=480 \
-          -i input_images=https://github.com/andreasjansson/sprites-dataset/raw/main/sprites.zip \
-          -i steps=100 \
-          ostris/flux-dev-lora-trainer
-
-    - name: Run cog-safe-push to test the predictor and optionally push to production
+    - name: Run cog-safe-push to test the trainer and predictor and optionally push to production
       env:
         ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
         REPLICATE_API_TOKEN: ${{ secrets.REPLICATE_API_TOKEN }}
       run: |
-        cog-safe-push -vv \
-          --test-model=replicate-internal/test-flux-fine-tuner \
-          ${{ github.event.inputs.test_only == 'true' && '--test-only' || '' }} \
-          ${{ github.event.inputs.compare_outputs == 'false' && '--no-compare-outputs' || '' }} \
-          -i replicate_weights="https://replicate.delivery/yhqm/iWjMZHd2T35kI5jaUkaG3Jb43MeA67PpYjKZQeifvTEf9yTNB/trained_model.tar" \
-          -i extra_lora="(omit)^50%" \
-          -i extra_lora="fofr/flux-pixar-cars" \
-          -i extra_lora="huggingface.co/fofr/flux-wrong" \
-          -i extra_lora="https://civitai.com/api/download/models/731747?type=Model&format=SafeTensor" \
-          --predict-timeout=600 \
-          ostris/flux-dev-lora-trainer
+        cog-safe-push -vv ${{ github.event.inputs.no_push == 'true' && '--no-push' || '' }}
diff --git a/cog-safe-push.yaml b/cog-safe-push.yaml
@@ -0,0 +1,119 @@
+model: ostris/flux-dev-lora-trainer
+test_model: replicate-internal/test-flux-fine-tuner
+predict:
+  compare_outputs: false
+  predict_timeout: 300
+  test_cases:
+
+    # basic
+    - inputs:
+        prompt: PLSM style, an image of a tiger, made of yellow, red, and purple goo plasma
+        num_outputs: 1
+        lora_scale: 1.1
+        num_inference_steps: 28
+        guidance_scale: 1.5
+        seed: 5259
+        output_format: jpg
+        replicate_weights: https://replicate.delivery/yhqm/MHJmIF7zlLKXPNN9LKN5yxYUuC7SyzKjMBtqeVwUGdPCpqqJA/trained_model.tar
+      match_url: https://replicate.delivery/yhqm/1jbh5xu9KqaoAV0fG8z3rhaN6ei5fyQEHffTJePLUZdXi182E/out-0.jpg
+
+    # img2img and custom size
+    - inputs:
+        prompt: "PLSM style, an image of a formula one car, made of yellow, red, and purple goo plasma"
+        num_outputs: 1
+        lora_scale: 1.0
+        aspect_ratio: "custom"
+        width: 500
+        height: 256
+        num_inference_steps: 20
+        model: "dev"
+        guidance_scale: 1.5
+        replicate_weights: "https://replicate.delivery/yhqm/MHJmIF7zlLKXPNN9LKN5yxYUuC7SyzKjMBtqeVwUGdPCpqqJA/trained_model.tar"
+        output_format: "jpg"
+        output_quality: 90
+        image: "https://storage.googleapis.com/cog-safe-push-public/fast-car.jpg"
+        prompt_strength: 0.9
+        seed: 12919
+      match_url: "https://replicate.delivery/yhqm/ySRQ9Kchm3pQAZmG5CsWuJJnrIUmI7KfffURlgNXRccTBo3mA/out-0.jpg"
+
+    # inpainting
+    - inputs:
+        prompt: "PLSM style, an image of a tiger, made of yellow, red, and purple goo plasma"
+        num_outputs: 1
+        lora_scale: 0.8
+        aspect_ratio: custom
+        width: 500
+        height: 256
+        num_inference_steps: 20
+        model: "dev"
+        guidance_scale: 1.5
+        replicate_weights: "https://replicate.delivery/yhqm/MHJmIF7zlLKXPNN9LKN5yxYUuC7SyzKjMBtqeVwUGdPCpqqJA/trained_model.tar"
+        output_format: "jpg"
+        output_quality: 90
+        prompt_strength: 1.0
+        image: "https://replicate.delivery/pbxt/LbXcDNOFtOMiBChe9QHHbeNJ5KcoTNfWEGmIxswhaexvXWEs/Screenshot%202024-09-11%20at%2014.36.07.png"
+        mask: "https://replicate.delivery/pbxt/LbXcE8bXqlTqPWIDb3SFFuJzIUMKSUscnqF97XVv0K2coEwr/Screenshot%202024-09-11%20at%2014.37.41.png"
+        seed: 64090
+      match_url: "https://replicate.delivery/yhqm/NwelDMqsy234YaQesvcpWxsxmlKYlsgRkBW9JBXAfomXYo3mA/out-0.jpg"
+
+    # aspect ratio
+    - inputs:
+        prompt: PLSM style, an image of a tiger, made of yellow, red, and purple goo plasma
+        aspect_ratio: "16:9"
+        num_outputs: 1
+        lora_scale: 1.1
+        num_inference_steps: 28
+        guidance_scale: 1.5
+        seed: 16726
+        output_format: jpg
+        replicate_weights: https://replicate.delivery/yhqm/MHJmIF7zlLKXPNN9LKN5yxYUuC7SyzKjMBtqeVwUGdPCpqqJA/trained_model.tar
+      match_url: https://replicate.delivery/yhqm/sedcOpLQ48XWE6KrjhRL9sIkuH6AJJI6Xix23ArqZ2wNr5tJA/out-0.jpg
+
+    # schnell
+    - inputs:
+        prompt: PLSM style, an image of a tiger, made of yellow, red, and purple goo plasma
+        num_outputs: 1
+        lora_scale: 1.1
+        num_inference_steps: 4
+        model: schnell
+        guidance_scale: 1.5
+        seed: 16726
+        output_format: jpg
+        replicate_weights: https://replicate.delivery/yhqm/MHJmIF7zlLKXPNN9LKN5yxYUuC7SyzKjMBtqeVwUGdPCpqqJA/trained_model.tar
+      match_url: https://replicate.delivery/yhqm/hHjTKpqSfjy4CqNzmNoWcYGW4WkYH069E4GcMvlkjaSRr5tJA/out-0.jpg
+
+    # extra lora
+    - inputs:
+        prompt: PLSM style, an image of a man ANDRS, made of yellow, red, and purple goo plasma
+        num_outputs: 1
+        lora_scale: 0.8
+        num_inference_steps: 20
+        model: dev
+        guidance_scale: 1.5
+        seed: 413
+        extra_lora: andreasjansson/flux-me
+        extra_lora_scale: 1.1
+        replicate_weights: https://replicate.delivery/yhqm/MHJmIF7zlLKXPNN9LKN5yxYUuC7SyzKjMBtqeVwUGdPCpqqJA/trained_model.tar
+        output_format: jpg
+        output_quality: 90
+      match_url: https://replicate.delivery/yhqm/jpyztFCfmetarE6p9Zt8VL2wfJpzyssOEtrX3ieGLspEee82E/out-0.jpg
+
+  fuzz:
+    fixed_inputs:
+      replicate_weights: https://replicate.delivery/yhqm/MHJmIF7zlLKXPNN9LKN5yxYUuC7SyzKjMBtqeVwUGdPCpqqJA/trained_model.tar
+    disabled_inputs:
+      - extra_lora
+    duration: 120
+    iterations: 5
+
+train:
+  train_timeout: 600
+  test_cases:
+
+    # tiny training with autocaption
+    - inputs:
+        input_images: "https://replicate.delivery/pbxt/LbXqvBzG3Bfo6CDkIkBNo8a8v6GWZGQrqxBrgeRpaQ1knkyD/me-tiny.zip"
+        steps: 50
+        autocaption: true
+        trigger_word: NDRS
+      match_prompt: "A dictionary of a version and weights which are a .tar file"