ideogram_block.py

import os
import requests
from typing import Optional
from PIL import Image
from io import BytesIO
from diffusers.utils import load_image
from diffusers.modular_pipelines import (
    PipelineState,
    ModularPipelineBlocks,
    InputParam,
    OutputParam,
)

IDEOGRAM_GENERATE_URL = "https://api.ideogram.ai/v1/ideogram-v3/generate"

class CreateCharacterImageBlock(ModularPipelineBlocks):
    
    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                name="prompt",
                type_hint=str,
                description="Text prompt describing the desired image."
            ),
            InputParam(
                name="character_image",
                type_hint=Image.Image,
                description="A single PIL Image to use as the character reference."
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="image",
                type_hint=Image.Image,
                description="Generated image."
            ),
            OutputParam(
                name="image_url",
                type_hint=str,
                description="URL to the generated image."
            )
        ]
    
    @staticmethod
    def generate_image_with_character_reference(
        prompt: str,
        character_image: Image.Image,        *,
        style_type: str = "AUTO",  # "GENERAL", "REALISTIC", "DESIGN", "FICTION"
        seed: Optional[int] = None,
        resolution: Optional[str] = "1024x1024",
        request_timeout: int = 60,
    ) -> str:
        """
        Generate an image using Ideogram's character reference feature.

        Args:
            api_key: Your Ideogram API key.
            prompt: Text prompt describing the desired image.
            character_image: A single PIL Image to use as the character reference.
            style_type: Ideogram style selection. "AUTO" lets the API decide.
            seed: Random seed for reproducibility.
            resolution: Image resolution.
            request_timeout: Timeout for HTTP requests (seconds).

        Returns:
            str: URL to the generated image (default)

        Raises:
            RuntimeError on API errors or missing results.
        """
        # Resolve API key from environment
        api_key = os.getenv("IDEOGRAM_API_KEY")

        if not api_key:
            raise RuntimeError(
                "IDEOGRAM_API_KEY is not set. Provide api_key param or set env var IDEOGRAM_API_KEY."
            )

        headers = {"Api-Key": api_key}

        # Validate image type
        if not isinstance(character_image, Image.Image):
            raise TypeError("character_image must be a PIL.Image.Image")

        # Validate seed and mutually exclusive resolution/aspect_ratio
        if seed is not None:
            if not (0 <= int(seed) <= 2147483647):
                raise ValueError("seed must be between 0 and 2147483647 inclusive")

        # Build multipart with a single character reference image
        with BytesIO() as buf:
            character_image.save(buf, format="PNG")
            buf.seek(0)
            files = [
                (
                    "character_reference_images",
                    ("character.png", buf, "image/png"),
                )
            ]

            data: dict = {
                "prompt": prompt,
                "style_type": style_type,
            }
            if seed is not None:
                data["seed"] = str(int(seed))
            if resolution is not None:
                data["resolution"] = resolution

            resp = requests.post(
                IDEOGRAM_GENERATE_URL,
                headers=headers,
                data=data,
                files=files,
                timeout=request_timeout,
            )
            if not resp.ok:
                raise RuntimeError(f"Ideogram API error {resp.status_code}: {resp.text}")

            payload = resp.json()
            # Expected: payload['data'][0]['url']
            try:
                image_url = payload["data"][0]["url"]
            except Exception as e:
                raise RuntimeError(f"Unexpected Ideogram response format: {payload}") from e

            return image_url

    def __call__(self, components: dict, state: PipelineState):
        
        block_state = self.get_block_state(state)

        if isinstance(block_state.character_image, str):
            character_image = load_image(block_state.character_image)
        elif isinstance(block_state.character_image, Image.Image):
            character_image = block_state.character_image
        else:
            raise ValueError(f"Invalid character image type: {type(block_state.character_image)}")

        block_state.image_url = self.generate_image_with_character_reference(
            prompt=block_state.prompt,
            character_image=character_image,
        )
        block_state.image = load_image(block_state.image_url)
        self.set_block_state(state, block_state)
        return components, state