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 (7debafd) with develop (e866345)}

Summary

✅ 1 untouched benchmarks

[gabrielmbmb]

Very cool! I think we need to fix some issues related to inheritance, but maybe we can tackle those in a separate PRs before the release.

[gabrielmbmb]

The issue here is that when calling super().get_runtime_parameters_info(), the InferenceEndpointsLLM.get_runtime_parameters_info method is being called instead of the one from RuntimeParametersMixin class, which is returning the dictionary without the optional key (already popped). I guess similar it's happening with OpenAI class.

[gabrielmbmb]

I think the inheritance from InferenceEndpointsLLM is a bit messy, not only in this method but also in other parts of the class. For example:

from distilabel.models import InferenceEndpointsImageGeneration, InferenceEndpointsLLM

igm = InferenceEndpointsImageGeneration()

This should raise a ValidationError because of the validator only_one_of_model_id_endpoint_name_or_base_url_provided (no endpoint name or model id), but instead gives TypeError:

TypeError: ValidationError.__new__() missing 1 required positional argument: 
'line_errors'

which is caused by the multiple inheritance I think.

[gabrielmbmb]

I think we should:

1. Move extended methods from RuntimeParametersMixins like runtime_parameters_names, get_runtime_parameters_info and generate_parsed_docstring to another class that then we can use in the base LLM, ImageGenerationModel and the future base classes to come.
2. Just duplicate the code for classes that uses a client like inference endpoints or openai, because inheriting from the LLM class it's messy or can get messy. Maybe we can create a base class that offers the client functionality for openai, inference endpoints, etc and then use it in the LLMs, ImageGenerationModel, etc.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[gabrielmbmb]

LGTM!