Create ideogram_block.py

ideogram_block.py

@@ -0,0 +1,150 @@
+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