first commit

jailbreak_toolbox/models/base_image_generator.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Union, Any
+from PIL import Image
+
+class BaseImageGenerator(ABC):
+    """
+    Abstract base class for image generation models.
+    Defines the interface that all image generation models must implement.
+    """
+    def __init__(self, **kwargs):
+        """Initialize with model-specific parameters."""
+        pass
+
+    @abstractmethod
+    def generate(self, prompt: Any, output_path: Optional[str] = None) -> Image.Image:
+        """
+        Generate an image from the given prompt.
+        
+        Args:
+            prompt: Input for image generation (can be text, image, or other data)
+            output_path: Optional path to save the generated image
+            
+        Returns:
+            PIL.Image: Generated image
+        """
+        pass
+
+    def save_image(self, image: Image.Image, output_path: str) -> None:
+        """
+        Helper method to save the generated image.
+        
+        Args:
+            image: PIL Image to save
+            output_path: Path where to save the image
+        """
+        image.save(output_path)