Add inference helpers & tests

.github/workflows/CODEOWNERS

@@ -0,0 +1 @@
+.github @Stability-AI/infrastructure

.github/workflows/test-inference.yml

@@ -0,0 +1,34 @@
+name: Test inference
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  test:
+    name: "Test inference"
+    # This action is designed only to run on the Stability research cluster at this time, so many assumptions are made about the environment
+    if: github.repository == 'stability-ai/generative-models'
+    runs-on: [self-hosted, slurm, g40]
+    steps:
+      - uses: actions/checkout@v3
+      - name: "Symlink checkpoints"
+        run: ln -s ${{secrets.SGM_CHECKPOINTS_PATH}} checkpoints
+      - name: "Setup python"
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+      - name: "Install Hatch"
+        run: pip install hatch
+      - name: "Run inference tests"
+        run: hatch run ci:test-inference --junit-xml test-results.xml
+      - name: Surface failing tests
+        if: always()
+        uses: pmeier/pytest-results-action@main
+        with:
+          path: test-results.xml
+          summary: true
+          display-options: fEX
+          fail-on-empty: true

pyproject.toml

@@ -32,3 +32,19 @@ include = [
 
 [tool.hatch.build.targets.wheel.force-include]
 "./configs" = "sgm/configs"
+
+[tool.hatch.envs.ci]
+# Skip for now, since requirements.txt is used by scripts and includes the project
+# This should be changed when dependencies are handled by Hatch
+skip-install = true 
+
+dependencies = [
+    "pytest"
+]
+
+[tool.hatch.envs.ci.scripts]
+test-inference = [
+    "pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2+cu118 --index-url https://download.pytorch.org/whl/cu118",
+    "pip install -r requirements_pt2.txt",
+    "pytest -v tests/inference/test_inference.py {args}",
+]

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+markers = 
+  inference: mark as inference test (deselect with '-m "not inference"')

requirements_pt13.txt

@@ -13,6 +13,7 @@ torchvision==0.14.1+cu117
 torchmetrics
 opencv-python==4.6.0.66
 fairscale
+pydantic
 pytorch-lightning==1.8.5
 fsspec
 kornia==0.6.9

requirements_pt2.txt

@@ -13,6 +13,7 @@ torchmetrics
 torchvision>=0.15.2
 opencv-python==4.6.0.66
 fairscale
+pydantic
 pytorch-lightning==2.0.1
 fire
 fsspec

scripts/demo/sampling.py

@@ -1,6 +1,13 @@
+import numpy as np
 from pytorch_lightning import seed_everything
 from scripts.demo.streamlit_helpers import *
 from scripts.util.detection.nsfw_and_watermark_dectection import DeepFloydDataFiltering
+from sgm.inference.helpers import (
+    do_img2img,
+    do_sample,
+    get_unique_embedder_keys_from_conditioner,
+    perform_save_locally,    
+)
 
 SAVE_PATH = "outputs/demo/txt2img/"
 
@@ -130,6 +137,8 @@ def run_txt2img(
 
     if st.button("Sample"):
         st.write(f"**Model I:** {version}")
+        outputs = st.empty()
+        st.text("Sampling")
         out = do_sample(
             state["model"],
             sampler,
@@ -143,6 +152,8 @@ def run_txt2img(
             return_latents=return_latents,
             filter=filter,
         )
+        show_samples(out, outputs)
+
         return out
 
 
@@ -174,6 +185,8 @@ def run_img2img(
     num_samples = num_rows * num_cols
 
     if st.button("Sample"):
+        outputs = st.empty()
+        st.text("Sampling")
         out = do_img2img(
             repeat(img, "1 ... -> n ...", n=num_samples),
             state["model"],
@@ -183,7 +196,9 @@ def run_img2img(
             force_uc_zero_embeddings=["txt"] if not is_legacy else [],
             return_latents=return_latents,
             filter=filter,
-        )
+            logger=st,
+        )        
+        show_samples(out, outputs)
         return out
 
 
@@ -248,8 +263,6 @@ def apply_refiner(
     save_locally, save_path = init_save_locally(os.path.join(SAVE_PATH, version))
 
     state = init_st(version_dict)
-    if state["msg"]:
-        st.info(state["msg"])
     model = state["model"]
 
     is_legacy = version_dict["is_legacy"]
@@ -274,7 +287,6 @@ def apply_refiner(
 
         version_dict2 = VERSION2SPECS[version2]
         state2 = init_st(version_dict2)
-        st.info(state2["msg"])
 
         stage2strength = st.number_input(
             "**Refinement strength**", value=0.3, min_value=0.0, max_value=1.0
@@ -314,6 +326,7 @@ def apply_refiner(
         samples_z = None
 
     if add_pipeline and samples_z is not None:
+        outputs = st.empty()
         st.write("**Running Refinement Stage**")
         samples = apply_refiner(
             samples_z,
@@ -323,7 +336,8 @@ def apply_refiner(
             prompt=prompt,
             negative_prompt=negative_prompt if is_legacy else "",
             filter=filter,
-        )
+        )    
+        show_samples(samples, outputs)    
 
     if save_locally and samples is not None:
         perform_save_locally(save_path, samples)