1 files changed, 1 insertions, 1 deletions

diff --git a/g4f/Provider/needs_auth/DeepInfraImage.py b/g4f/Provider/needs_auth/DeepInfraImage.py
index 2310c1c8..24df04e3 100644
--- a/g4f/Provider/needs_auth/DeepInfraImage.py
+++ b/g4f/Provider/needs_auth/DeepInfraImage.py
@@ -73,7 +73,7 @@ class DeepInfraImage(AsyncGeneratorProvider, ProviderModelMixin):
             async with session.post(f"{api_base.rstrip('/')}/{model}", json=data) as response:
                 await raise_for_status(response)
                 data = await response.json()
-                images = data["output"] if "output" in data else data["images"]
+                images = data.get("output", data.get("images"))
                 if not images:
                     raise RuntimeError(f"Response: {data}")
                 images = images[0] if len(images) == 1 else images