Upscale (#186)

* Adding upscaler functionality * Update help and readme * Update notebook with upscaler * Specs and test * Client side message size limits increased * Update README.md * Bump version
Stability-AI · Mar 13, 2023 · f01a647 · f01a647
1 parent 2fab18d
commit f01a647
Show file tree

Hide file tree

Showing 6 changed files with 321 additions and 62 deletions.
diff --git a/README.md b/README.md
@@ -25,18 +25,21 @@ You can manage API keys in your dreamstudio account [here](https://beta.dreamstu
 
 Then to invoke:
 
-`python3 -m stability_sdk.client -W 512 -H 512 "A stunning house."`
+`python3 -m stability_sdk generate -W 512 -H 512 "A stunning house."`
 
 It will generate and put PNGs in your current directory.
 
+To upscale:
+`python3 -m stability_sdk upscale -i "/path/to/image.png"`
+
 ## SDK Usage
 
 See usage demo notebooks in ./nbs
 
 ## Command line usage
 
 ```
-usage: python -m stability_sdk [-h] [--height HEIGHT] [--width WIDTH] [--start_schedule START_SCHEDULE]
+usage: python -m stability_sdk generate [-h] [--height HEIGHT] [--width WIDTH] [--start_schedule START_SCHEDULE]
                  [--end_schedule END_SCHEDULE] [--cfg_scale CFG_SCALE] [--sampler SAMPLER]
                  [--steps STEPS] [--seed SEED] [--prefix PREFIX] [--engine ENGINE]
                  [--num_samples NUM_SAMPLES] [--artifact_types ARTIFACT_TYPES]
@@ -80,6 +83,31 @@ options:
   --mask_image MASK_IMAGE, -m MASK_IMAGE
                         Mask image
 ```
+For upscale:
+```
+usage: client.py upscale
+       [-h]
+       --init_image INIT_IMAGE
+       [--height HEIGHT] [--width WIDTH] [--prefix PREFIX] [--artifact_types ARTIFACT_TYPES]
+       [--no-store] [--show] [--engine ENGINE]
+
+options:
+  -h, --help            show this help message and exit
+  --init_image INIT_IMAGE, -i INIT_IMAGE
+                        Init image
+  --height HEIGHT, -H HEIGHT
+                        height of upscaled image in pixels
+  --width WIDTH, -W WIDTH
+                        width of upscaled image in pixels
+  --prefix PREFIX, -p PREFIX
+                        output prefixes for artifacts
+  --artifact_types ARTIFACT_TYPES, -t ARTIFACT_TYPES
+                        filter artifacts by type (ARTIFACT_IMAGE, ARTIFACT_TEXT, ARTIFACT_CLASSIFICATIONS, etc)
+  --no-store            do not write out artifacts
+  --show                open artifacts using PIL
+  --engine ENGINE, -e ENGINE
+                        engine to use for upscale
+```
 
 
 ## Connecting to the API using languages other than Python

diff --git a/nbs/demo_colab.ipynb b/nbs/demo_colab.ipynb
@@ -300,11 +300,33 @@
     "            print('GUIDANCE: SLOWER:')\n",
     "            display(img5)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Upscaling\n",
+    "answers = stability_api.upscale(\n",
+    "    init_image=img3\n",
+    ")\n",
+    "\n",
+    "for resp in answers:\n",
+    "    for artifact in resp.artifacts:\n",
+    "        if artifact.finish_reason == generation.FILTER:\n",
+    "            warnings.warn(\n",
+    "                \"Your request activated the API's safety filters and could not be processed.\"\n",
+    "                \"Please submit a different image and try again.\")\n",
+    "        if artifact.type == generation.ARTIFACT_IMAGE:\n",
+    "            img3 = Image.open(io.BytesIO(artifact.binary))\n",
+    "            display(img3)"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.9.7 ('dmarx-je5LfYh2')",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -318,12 +340,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.7"
+   "version": "3.9.5"
   },
   "orig_nbformat": 4,
   "vscode": {
    "interpreter": {
-    "hash": "57881a85d677a34ea29564e0084ef84f4058c4e30a2bb466eb0e0b908d0628df"
+    "hash": "f9f85f796d01129d0dd105a088854619f454435301f6ffec2fea96ecbd9be4ac"
    }
   }
  },

diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name='stability-sdk',
-    version='0.3.2',
+    version='0.4.0',
     author='Wes Brown',
     author_email='[email protected]',
     maintainer='David Marx',