Examples

Nano-banana on CPU

from typing import List, Optional

from pydantic import BaseModel

from fma.toolkit import model as fma_model
from fma.toolkit.fields.image import Image


class Input(BaseModel):
    prompt: str
    image: Optional[Image]
    image1: Optional[Image]


class Output(BaseModel):
    sample: Image


class Model(fma_model.Model):
    requirements: List[str] = ["git+https://<token>@github.com/FlyMyAI/flymy_nano_banana.git"]

    def initialize(self):
        import os

        os.environ['GOOGLE_API_KEY'] = ''
        os.environ['MODEL'] = 'gemini-2.5-flash-image-preview'

        from flymy_nano_banana.image_generator import GoogleImageGenerator

        self.image_generator = GoogleImageGenerator()

    def predict(self, input: Input) -> Output:
        prompt = input["prompt"]
        image = input["image"]
        image1 = input["image1"]

        generated_image = self.image_generator.generate_image(prompt, image, image1)  # type: ignore
        if generated_image is None:
            raise ValueError("No image was generated")

        return {"sample": generated_image}