Ideogram 4

"""Ideogram 4, generated from a declarative template."""

from __future__ import annotations

import base64
import io

from pydantic import BaseModel, ConfigDict, Field

from runware_serverless import serve

MODEL = "ideogram-ai/ideogram-4-nf4-diffusers"
REVISION = "1874bc70267ba2c823a7239e1d70dd308c8d64dc"
CACHE = "/weights"
DEFAULT_STEPS = 48
GUIDANCE = None
"""A distilled model is trained with guidance at 1, so the pipeline's own
default would push it somewhere it was never trained to go. None leaves the
pipeline to decide."""
SIZE = {"width": 1024, "height": 1024}
"""Passed on every call because the pipeline's own default is this model's
maximum, and the cost is the square of it."""


class Request(BaseModel):
    # An input this model does not read is a line the caller believed was doing
    # something, so it is refused rather than dropped. Without this, pydantic
    # ignores it: a width sent to a model of fixed size, or a misspelled seed,
    # both come back as a result that silently answered a different request.
    model_config = ConfigDict(extra="forbid")

    prompt: str = Field(description="What to generate.")
    seed: int | None = Field(default=None, description="Repeat a result by repeating its seed.")


class Generated(BaseModel):
    image: str = Field(description="The result as a base64 PNG.")
    steps: int
    seed: int | None


@serve
class GeneratedImage:
    def load(self) -> None:
        if hasattr(self, "pipeline"):
            return

        # Imported here rather than at the top of the file. The build imports
        # this file with the platform's own dependencies, so nothing above
        # exists until a worker runs.
        import torch
        from diffusers import DiffusionPipeline

        self.torch = torch
        self.pipeline = DiffusionPipeline.from_pretrained(
            MODEL,
            revision=REVISION,
            cache_dir=CACHE,
            # Each component lands on the card as it loads, rather than the
            # pipeline being assembled on the host and moved. A worker's host RAM
            # is half its card's VRAM, so moving a built pipeline has a ceiling
            # well below the card's, and the kill arrives as the container's exit
            # 137 rather than as a CUDA error. Moving one also refuses weights
            # that arrive on the meta device, which a pre-quantized checkpoint's
            # do.
            device_map="cuda",
        )

        # The first run compiles kernels and allocates, so no request pays for
        # it. It also moves a failure into the readiness check, where it is
        # reported, rather than into somebody's first call.
        #
        # At the entry's own step count rather than at one. A pipeline carrying
        # a guidance schedule requires the two to agree, and refuses a single
        # step with "guidance_schedule must have length num_inference_steps".
        # Warming at the real count also exercises the path a request takes.
        self.pipeline(prompt="warmup", num_inference_steps=DEFAULT_STEPS, **SIZE)

    def predict(self, request: Request) -> Generated:
        generator = None
        if request.seed is not None:
            generator = self.torch.Generator(device="cuda").manual_seed(request.seed)

        extra = {} if GUIDANCE is None else {"guidance_scale": GUIDANCE}

        image = self.pipeline(
            prompt=request.prompt,
            num_inference_steps=DEFAULT_STEPS,
            generator=generator,
            **SIZE,
            **extra,
        ).images[0]

        buffer = io.BytesIO()
        image.save(buffer, format="PNG")

        return Generated(
            image=base64.b64encode(buffer.getvalue()).decode("ascii"),
            steps=DEFAULT_STEPS,
            seed=request.seed,
        )

configVersion: 1
id: ideogram-4
name: Ideogram 4
source:
  type: code
  entry: model.py
compute:
  gpuType: l40s
build:
  requirements:
    - torch==2.14.1
    - torchvision==0.29.1
    - diffusers==0.42.1
    - transformers==5.21.0
    - accelerate==1.15.0
    - bitsandbytes==0.50.2
  baseImage: python:3.12
volumes:
  - /weights
secrets:
  - name: HF_TOKEN