feat(bedrock.py): add stable diffusion image generation support

2025-04-26 03:04:13 +00:00 · 2024-02-03 12:08:38 -08:00 · 2024-02-03 12:08:38 -08:00 · 36416360c4
commit 36416360c4
parent 30c96ee872
6 changed files with 314 additions and 17 deletions
--- a/litellm/main.py
+++ b/litellm/main.py
@ -3153,7 +3153,18 @@ def image_generation(
                model_response=model_response,
                aimg_generation=aimg_generation,
            )
-
+        elif custom_llm_provider == "bedrock":
+            if model is None:
+                raise Exception("Model needs to be set for bedrock")
+            model_response = bedrock.image_generation(
+                model=model,
+                prompt=prompt,
+                timeout=timeout,
+                logging_obj=litellm_logging_obj,
+                optional_params=optional_params,
+                model_response=model_response,
+                aimg_generation=aimg_generation,
+            )
        return model_response
    except Exception as e:
        ## Map to OpenAI Exception