Image Language Models and ImageGeneration task

[plaguss]

Description

This PR adds a new module to models: models/image_generation to store image models (InferenceEndpointsImageGeneration and OpenAIImageGeneration), with 2 new base classes: ImageGenerationModel and AsyncImageGenerationModel, and a new ImageGeneration task.

Sample pipeline and dataset. Take into account the distiset.transform_columns_to_image method, necessary to push the dataset with the images as objects instead of strings.

from datasets import load_dataset

from distilabel.models.image_generation import InferenceEndpointsImageGeneration
from distilabel.pipeline import Pipeline
from distilabel.steps import KeepColumns
from distilabel.steps.tasks import ImageGeneration

ds = load_dataset("dvilasuero/finepersonas-v0.1-tiny", split="train").select(range(3))

with Pipeline(name="image_generation_pipeline") as pipeline:
    igm = InferenceEndpointsImageGeneration(model_id="black-forest-labs/FLUX.1-schnell")

    img_generation = ImageGeneration(
        name="flux_schnell", image_generation_model=igm, input_mappings={"prompt": "persona"}
    )

    keep_columns = KeepColumns(columns=["persona", "model_name", "image"])

    img_generation >> keep_columns


if __name__ == "__main__":
    distiset = pipeline.run(use_cache=False, dataset=ds)
    # Save the images as `PIL.Image.Image`
    distiset = distiset.transform_columns_to_image("image")
    distiset.push_to_hub("plaguss/test-finepersonas-v0.1-tiny-flux-schnell")

[github-actions]

Documentation for this PR has been built. You can view it at: https://distilabel.argilla.io/pr-1060/

[codspeed-hq]

CodSpeed Performance Report

Merging #1060 will not alter performance

_{Comparing vision-language-models (e9e6790) with develop (a8d02c2)}

Summary

✅ 1 untouched benchmarks

[davidberenstein1957]

@burtenshaw

[davidberenstein1957]

Suggested change

	+ distiset = distiset.transform_columns_to_image("image")
	distiset = distiset.transform_columns_to_image("image")

[davidberenstein1957]

I just go a pil error due to some imports

/Users/davidberenstein/Documents/programming/argilla/distilabel/.venv/lib/python3.11/site-packag │
│ es/distilabel/steps/tasks/text_generation_with_image.py:18 in <module>                           │
│                                                                                                  │
│    15 from typing import TYPE_CHECKING, Any, Literal, Union                                      │
│    16                                                                                            │
│    17 from jinja2 import Template                                                                │
│ ❱  18 from PIL import Image                                                                      │
│    19 from pydantic import Field                                                                 │
│    20                                                                                            │
│    21 from distilabel.steps.tasks.base import Task                                               │
│                                                                                                  │
│ ╭──────────── locals ────────────╮                                                               │
│ │       Literal = typing.Literal │                                                               │
│ │ TYPE_CHECKING = False          │                                                               │
│ │         Union = typing.Union   │                                                               │
│ ╰────────────────────────────────╯                                                               │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ModuleNotFoundError: No module named 'PIL'