Files
EvoSynth/jailbreak_toolbox/models/base_image_generator.py
dongdongunique f3af94df0b first commit
2025-12-10 00:54:02 +08:00

37 lines
1.1 KiB
Python

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)