Updates error handling, image generation example notebook

jupyterlab · Apr 27, 2023 · 38a95a6 · 38a95a6
1 parent 45818d6
commit 38a95a6
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 5 deletions.
diff --git a/examples/images.ipynb b/examples/images.ipynb
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -191,6 +191,10 @@ def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         """
         _model_kwargs = self.model_kwargs or {}
         response = self.client(inputs=prompt, params=_model_kwargs)
+
+        if type(response) is dict and "error" in response:
+            raise ValueError(f"Error raised by inference API: {response['error']}")
+
         # Custom code for responding to image generation responses
         if self.client.task == "text-to-image":
             imageFormat = response.format # Presume it's a PIL ImageFile
@@ -209,8 +213,6 @@ def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
             # Encode image data to Base64 bytes, then decode bytes to str
             return (mimeType + ';base64,' + base64.b64encode(buffer.getvalue()).decode())
 
-        if "error" in response:
-            raise ValueError(f"Error raised by inference API: {response['error']}")
         if self.client.task == "text-generation":
             # Text generation return includes the starter text.
             text = response[0]["generated_text"][len(prompt) :]