Add image targets

Author: hupe1980

aisploit/core/__init__.py

@@ -7,7 +7,7 @@
 from .model import BaseChatModel, BaseEmbeddings, BaseLLM, BaseModel
 from .prompt import BasePromptValue
 from .report import BaseReport
-from .target import BaseTarget, Response
+from .target import BaseImageTarget, BaseTarget, ContentFilteredException, Response
 from .vectorstore import BaseVectorStore
 
 __all__ = [
@@ -30,6 +30,8 @@
     "BasePromptValue",
     "BaseReport",
     "BaseTarget",
+    "BaseImageTarget",
     "Response",
+    "ContentFilteredException",
     "BaseVectorStore",
 ]

aisploit/core/target.py

@@ -1,10 +1,18 @@
+import base64
+import io
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Dict
+from typing import Any, Dict, Literal
+
+from PIL import Image
 
 from .prompt import BasePromptValue
 
 
+class ContentFilteredException(Exception):
+    pass
+
+
 @dataclass
 class Response:
     """A class representing a response from the target."""
@@ -32,3 +40,15 @@ def send_prompt(self, prompt: BasePromptValue) -> Response:
             Response: The response from the target.
         """
         pass
+
+
+@dataclass
+class BaseImageTarget(ABC):
+    size: Literal["512x512", "1024x1024"] = "512x512"
+    show_image: bool = False
+
+    def _show_base64_image(self, base64_image: str) -> None:
+        base64_bytes = base64_image.encode("ascii")
+        image_bytes = base64.b64decode(base64_bytes)
+        image = Image.open(io.BytesIO(image_bytes))
+        image.show()

aisploit/targets/__init__.py

@@ -1,6 +1,6 @@
 from .copilot import CopilotClient, CopilotTarget
 from .email import EmailReceiver, EmailSender, EmailTarget, UserPasswordAuth
-from .image import OpenAIImageTarget
+from .image import BedrockAmazonImageTarget, BedrockStabilityImageTarget, OpenAIImageTarget
 from .langchain import LangchainTarget
 from .stdout import StdOutTarget
 from .target import WrapperTarget, target
@@ -12,6 +12,8 @@
     "EmailSender",
     "EmailReceiver",
     "UserPasswordAuth",
+    "BedrockAmazonImageTarget",
+    "BedrockStabilityImageTarget",
     "OpenAIImageTarget",
     "LangchainTarget",
     "StdOutTarget",

aisploit/targets/image.py

@@ -1,14 +1,117 @@
+import json
 import os
-from dataclasses import dataclass
+from abc import ABC
+from dataclasses import dataclass, field
 from typing import Optional
 
+import boto3
+from botocore.exceptions import ClientError
 from openai import OpenAI
 
-from ..core import BasePromptValue, BaseTarget, Response
+from ..core import BaseImageTarget, BasePromptValue, ContentFilteredException, Response
 
 
 @dataclass
-class OpenAIImageTarget(BaseTarget):
+class BaseBedrockImageTarget(BaseImageTarget, ABC):
+    session: boto3.Session = field(default_factory=lambda: boto3.Session())
+    region_name: str = "us-east-1"
+
+    def __post_init__(self):
+        self._client = self.session.client("bedrock-runtime", region_name=self.region_name)
+
+
+@dataclass
+class BedrockAmazonImageTarget(BaseBedrockImageTarget):
+    model: str = "titan-image-generator-v1"
+    quality: str = "standard"
+    seed: int = 0
+    cfg_scale: int = 8
+
+    def send_prompt(self, prompt: BasePromptValue) -> Response:
+        width, height = self.size.split("x")
+        body = {
+            "textToImageParams": {
+                "text": prompt.to_string(),
+            },
+            "taskType": "TEXT_IMAGE",
+            "imageGenerationConfig": {
+                "seed": self.seed,
+                "cfgScale": self.cfg_scale,
+                "quality": self.quality,
+                "width": int(width),
+                "height": int(height),
+                "numberOfImages": 1,
+            },
+        }
+
+        try:
+            response = self._client.invoke_model(
+                body=json.dumps(body),
+                modelId=f"amazon.{self.model}",
+            )
+
+            response_body = json.loads(response["body"].read())
+
+            if response_body["error"]:
+                raise Exception(response_body["error"])
+
+            base64_image = response_body["images"][0]
+
+            if self.show_image:
+                self._show_base64_image(base64_image)
+
+            return Response(content=base64_image)
+        except ClientError as e:
+            if e.response['Error']['Code'] == 'ValidationException':
+                if "blocked by our content filters" in e.response['Error']['Message']:
+                    raise ContentFilteredException(e.response['Error']['Message']) from e
+
+            raise e
+
+
+@dataclass
+class BedrockStabilityImageTarget(BaseBedrockImageTarget):
+    model: str = "stable-diffusion-xl-v1"
+    steps: int = 50
+    seed: int = 0
+    cfg_scale: int = 8
+
+    def send_prompt(self, prompt: BasePromptValue) -> Response:
+        width, height = self.size.split("x")
+        body = {
+            "text_prompts": [{"text": prompt.to_string(), "weight": 1}],
+            "seed": self.seed,
+            "cfg_scale": self.cfg_scale,
+            "width": int(width),
+            "height": int(height),
+            "steps": self.steps,
+        }
+
+        response = self._client.invoke_model(
+            body=json.dumps(body),
+            modelId=f"stability.{self.model}",
+        )
+
+        response_body = json.loads(response["body"].read())
+
+        finish_reason = response_body.get("artifacts")[0].get("finishReason")
+
+        if finish_reason == "CONTENT_FILTERED":
+            raise ContentFilteredException(f"Image error: {finish_reason}")
+
+        if finish_reason == "ERROR":
+            raise Exception(f"Image error: {finish_reason}")
+
+        base64_image = response_body["artifacts"][0]["base64"]
+
+        if self.show_image:
+            self._show_base64_image(base64_image)
+
+        return Response(content=base64_image)
+
+
+@dataclass
+class OpenAIImageTarget(BaseImageTarget):
     api_key: Optional[str] = None
 
     def __post_init__(self):
@@ -18,6 +121,16 @@ def __post_init__(self):
         self._client = OpenAI(api_key=self.api_key)
 
     def send_prompt(self, prompt: BasePromptValue) -> Response:
-        response = self._client.images.generate(prompt=prompt.to_string(), n=1)
-        print(response)
-        return Response(content="")
+        response = self._client.images.generate(
+            prompt=prompt.to_string(),
+            size=self.size,
+            n=1,
+            response_format="b64_json",
+        )
+
+        base64_image = response.data[0].b64_json
+
+        if self.show_image:
+            self._show_base64_image(base64_image)
+
+        return Response(content=base64_image)

examples/target.ipynb

@@ -49,6 +49,7 @@ python-docx = "^1.1.0"
 brotli = "^1.1.0"
 stdlib-list = "^0.10.0"
 presidio-analyzer = "^2.2.354"
+boto3 = "^1.34.88"
 
 [tool.poetry.group.dev.dependencies]
 chromadb = "^0.4.23"