Image Language Models and ImageGeneration task

docs/sections/how_to_guides/advanced/distiset.md

@@ -119,6 +119,33 @@ class MagpieGenerator(GeneratorTask, MagpieBase):
 
 The `Citations` section can include any number of bibtex references. To define them, you can add as much elements as needed just like in the example: each citation will be a block of the form: ` ```@misc{...}``` `. This information will be automatically used in the README of your `Distiset` if you decide to call `distiset.push_to_hub`. Alternatively, if the `Citations` is not found, but in the `References` there are found any urls pointing to `https://arxiv.org/`, we will try to obtain the `Bibtex` equivalent automatically. This way, Hugging Face can automatically track the paper for you and it's easier to find other datasets citing the same paper, or directly visiting the paper page.
 
+#### Image Datasets
+
+!!! info "Keep reading if you are interested in Image datasets"
+
+    The `Distiset` object has a new method `transform_columns_to_image` specifically to transform the images to `PIL.Image.Image` before pushing the dataset to the hugging face hub.
+
+Since version `1.5.0` we have the [`ImageGeneration`](https://distilabel.argilla.io/dev/components-gallery/task/imagegeneration/) task that is able to generate images from text. By default, all the process will work internally with a string representation for the images. This is done for simplicity while processing. But to take advantage of the huggingface hub functionalities if the dataset generate is going to be stored there, a proper Image object may be preferable, so we can see the images in the dataset viewer for example. Let's take a look at the following pipeline extracted from "examples/image_generation.py" at the root of the repository to see how we can do it:
+
+```diff
+# Assume all the imports are already done, we are only interested
+with Pipeline(name="image_generation_pipeline") as pipeline:
+    img_generation = ImageGeneration(
+        name="flux_schnell",
+        llm=igm,
+        InferenceEndpointsImageGeneration(model_id="black-forest-labs/FLUX.1-schnell")
+    )
+    ...
+
+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(...)
+```
+
+After calling [`transform_columns_to_image`][distilabel.distiset.Distiset.transform_columns_to_image] on the image columns we may have generated (in this case we only want to transform the `image` column, but a list can be passed). This will apply to any leaf nodes we have in the pipeline, meaning if we have different subsets, the "image" column will be found in all of them, or we can pass a list of columns.
+
 ### Save and load from disk
 
 Take into account that these methods work as `datasets.load_from_disk` and `datasets.Dataset.save_to_disk` so the arguments are directly passed to those methods. This means you can also make use of `storage_options` argument to save your [`Distiset`][distilabel.distiset.Distiset] in your cloud provider, including the distilabel artifacts (`pipeline.yaml`, `pipeline.log` and the `README.md` with the dataset card). You can read more in `datasets` documentation [here](https://huggingface.co/docs/datasets/filesystems#saving-serialized-datasets).

docs/sections/pipeline_samples/examples/image_generation.md

@@ -0,0 +1,108 @@
+---
+hide: toc
+---
+
+# Image generation with `distilabel`
+
+Create synthetic images using `distilabel`.
+
+This example shows how distilabel can be used to generate image data, either using [`InferenceEndpointsImageGeneration`](https://distilabel.argilla.io/dev/components-gallery/image_generation/inferenceendpointsimagegeneration/) or [`OpenAIImageGeneration`](https://distilabel.argilla.io/dev/components-gallery/image_generation/openaiimagegeneration/), thanks to the [`ImageGeneration`](https://distilabel.argilla.io/dev/components-gallery/task/imagegeneration/) task.
+
+
+=== "Inference Endpoints - black-forest-labs/FLUX.1-schnell"
+
+    ```python
+    from distilabel.pipeline import Pipeline
+    from distilabel.steps import KeepColumns
+    from distilabel.models.image_generation import InferenceEndpointsImageGeneration
+    from distilabel.steps.tasks import ImageGeneration
+
+    from datasets import load_dataset
+
+    ds = load_dataset("dvilasuero/finepersonas-v0.1-tiny", split="train").select(range(3))
+
+    with Pipeline(name="image_generation_pipeline") as pipeline:
+        ilm = InferenceEndpointsImageGeneration(
+            model_id="black-forest-labs/FLUX.1-schnell"
+        )
+
+        img_generation = ImageGeneration(
+            name="flux_schnell",
+            llm=ilm,
+            input_mappings={"prompt": "persona"}
+        )
+
+        keep_columns = KeepColumns(columns=["persona", "model_name", "image"])
+
+        img_generation >> keep_columns
+    ```
+
+    Sample image for the prompt:
+
+    > A local art historian and museum professional interested in 19th-century American art and the local cultural heritage of Cincinnati.
+
+    ![image_ie](https://huggingface.co/datasets/plaguss/test-finepersonas-v0.1-tiny-flux-schnell/resolve/main/artifacts/flux_schnell/images/3333f9870feda32a449994017eb72675.jpeg)
+
+=== "OpenAI - dall-e-3"
+
+    ```python
+    from distilabel.pipeline import Pipeline
+    from distilabel.steps import KeepColumns
+    from distilabel.models.image_generation import OpenAIImageGeneration
+    from distilabel.steps.tasks import ImageGeneration
+
+    from datasets import load_dataset
+
+    ds = load_dataset("dvilasuero/finepersonas-v0.1-tiny", split="train").select(range(3))
+
+    with Pipeline(name="image_generation_pipeline") as pipeline:
+        ilm = OpenAIImageGeneration(
+            model="dall-e-3",
+            generation_kwargs={
+                "size": "1024x1024",
+                "quality": "standard",
+                "style": "natural"
+            }
+        )
+
+        img_generation = ImageGeneration(
+            name="dalle-3"
+            llm=ilm,
+            input_mappings={"prompt": "persona"}
+        )
+
+        keep_columns = KeepColumns(columns=["persona", "model_name", "image"])
+
+        img_generation >> keep_columns
+    ```
+
+    Sample image for the prompt:
+
+    > A local art historian and museum professional interested in 19th-century American art and the local cultural heritage of Cincinnati.
+
+    ![image_oai](https://huggingface.co/datasets/plaguss/test-finepersonas-v0.1-tiny-dall-e-3/resolve/main/artifacts/dalle-3/images/3333f9870feda32a449994017eb72675.jpeg)
+
+!!! success "Save the Distiset as an Image Dataset"
+
+    Note the call to `Distiset.transform_columns_to_image`, to have the images uploaded directly as an [`Image dataset`](https://huggingface.co/docs/hub/en/datasets-image):
+
+    ```python
+    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")
+
+    ```
+
+The full pipeline can be run at the following example. Keep in mind, you need to install `pillow` first: `pip install distilabel[vision]`.
+
+??? Run
+
+    ```python
+    python examples/image_generation.py
+    ```
+
+```python title="image_generation.py"
+--8<-- "examples/image_generation.py"
+```

docs/sections/pipeline_samples/index.md

@@ -153,6 +153,14 @@ hide: toc
 
     [:octicons-arrow-right-24: Example](examples/exam_questions.md)
 
+-   __Image generation with distilabel__
+
+    ---
+
+    Generate synthetic images using distilabel.
+
+    [:octicons-arrow-right-24: Example](examples/image_generation.md)
+
 
 </div>
 

examples/image_generation.py

@@ -0,0 +1,42 @@
+# Copyright 2023-present, Argilla, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+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")

mkdocs.yml

@@ -218,6 +218,7 @@ nav:
           - Structured generation with instructor: "sections/pipeline_samples/examples/mistralai_with_instructor.md"
           - Create a social network with FinePersonas: "sections/pipeline_samples/examples/fine_personas_social_network.md"
           - Create questions and answers for a exam: "sections/pipeline_samples/examples/exam_questions.md"
+          - Image generation with distilabel: "sections/pipeline_samples/examples/image_generation.md"
   - API Reference:
       - Step:
           - "api/step/index.md"

pyproject.toml

@@ -102,6 +102,7 @@ text-clustering = [
     "scikit-learn >= 1.4.1",
     "matplotlib >= 3.8.3",   # For the figure (even though it's optional)
 ]
+vision = ["Pillow >= 10.3.0"]  # To work with images.
 
 # minhash
 minhash = ["datasketch >= 1.6.5", "nltk>3.8.1"]

src/distilabel/distiset.py

@@ -24,7 +24,7 @@
 
 import fsspec
 import yaml
-from datasets import Dataset, load_dataset, load_from_disk
+from datasets import Dataset, DatasetDict, load_dataset, load_from_disk
 from datasets.filesystems import is_remote_filesystem
 from huggingface_hub import DatasetCardData, HfApi, upload_file, upload_folder
 from huggingface_hub.file_download import hf_hub_download
@@ -184,9 +184,17 @@ def _get_card(
         """
         sample_records = {}
         for name, dataset in self.items():
-            sample_records[name] = (
+            sample_record = (
                 dataset[0] if not isinstance(dataset, dict) else dataset["train"][0]
             )
+            from PIL import ImageFile
+
+            for k, v in sample_record.items():
+                if isinstance(v, ImageFile.ImageFile):
+                    v = ""
+                sample_record[k] = v
+
+            sample_records[name] = sample_record
 
         readme_metadata = {}
         if repo_id and token:
@@ -574,6 +582,51 @@ def __repr__(self):
         repr = re.sub(r"^", " " * 4, repr, count=0, flags=re.M)
         return f"Distiset({{\n{repr}\n}})"
 
+    def transform_columns_to_image(self, columns: Union[str, list[str]]) -> Self:
+        """Transforms the columns of the dataset to `PIL.Image` objects.
+
+        Args:
+            columns: Column or list of columns to transform.
+
+        Returns:
+            Transforms the columns of the dataset to `PIL.Image` objects before pushing,
+            so the Hub treats them as Image objects and can be rendered in the dataset
+            viewer, and cast them to be automatically transformed when downloading
+            the dataset back.
+        """
+        from datasets import Image
+
+        from distilabel.models.image_generation.utils import image_from_str
+
+        columns = [columns] if isinstance(columns, str) else columns
+
+        def cast_to_image(row: dict) -> dict:
+            for column in columns:
+                row[column] = image_from_str(row[column])
+            return row
+
+        for name, dataset in self.items():
+            # In case train_test_split was called
+            if isinstance(dataset, DatasetDict):
+                for split, dataset_split in dataset.items():
+                    dataset_split = dataset_split.map(cast_to_image)
+                    for column in columns:
+                        if column in dataset_split.column_names:
+                            dataset_split = dataset_split.cast_column(
+                                column, Image(decode=True)
+                            )
+                    self[name][split] = dataset_split
+            else:
+                dataset = dataset.map(cast_to_image)
+
+                for column in columns:
+                    if column in dataset.column_names:
+                        dataset = dataset.cast_column(column, Image(decode=True))
+
+                self[name] = dataset
+
+        return self
+
 
 def create_distiset(  # noqa: C901
     data_dir: Path,

src/distilabel/models/__init__.py

@@ -18,6 +18,14 @@
     SentenceTransformerEmbeddings,
 )
 from distilabel.models.embeddings.vllm import vLLMEmbeddings
+from distilabel.models.image_generation.base import (
+    AsyncImageGenerationModel,
+    ImageGenerationModel,
+)
+from distilabel.models.image_generation.huggingface.inference_endpoints import (
+    InferenceEndpointsImageGeneration,
+)
+from distilabel.models.image_generation.openai import OpenAIImageGeneration
 from distilabel.models.llms.anthropic import AnthropicLLM
 from distilabel.models.llms.anyscale import AnyscaleLLM
 from distilabel.models.llms.azure import AzureOpenAILLM
@@ -46,13 +54,17 @@
     "CohereLLM",
     "GroqLLM",
     "InferenceEndpointsLLM",
+    "ImageGenerationModel",
+    "AsyncImageGenerationModel",
+    "InferenceEndpointsImageGeneration",
     "LiteLLM",
     "LlamaCppLLM",
     "MistralLLM",
     "CudaDevicePlacementMixin",
     "MixtureOfAgentsLLM",
     "OllamaLLM",
     "OpenAILLM",
+    "OpenAIImageGeneration",
     "TogetherLLM",
     "TransformersLLM",
     "GenerateOutput",

src/distilabel/models/embeddings/base.py

@@ -50,7 +50,9 @@ class Embeddings(RuntimeParametersMixin, BaseModel, _Serializable, ABC):
 
     def load(self) -> None:
         """Method to be called to initialize the `Embeddings`"""
-        self._logger = logging.getLogger(f"distilabel.llm.{self.model_name}")
+        self._logger = logging.getLogger(
+            f"distilabel.models.embeddings.{self.model_name}"
+        )
 
     def unload(self) -> None:
         """Method to be called to unload the `Embeddings` and release any resources."""

src/distilabel/models/image_generation/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2023-present, Argilla, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from distilabel.models.image_generation.base import (
+    AsyncImageGenerationModel,
+    ImageGenerationModel,
+)
+from distilabel.models.image_generation.huggingface.inference_endpoints import (
+    InferenceEndpointsImageGeneration,
+)
+from distilabel.models.image_generation.openai import OpenAIImageGeneration
+
+__all__ = [
+    "ImageGenerationModel",
+    "AsyncImageGenerationModel",
+    "InferenceEndpointsImageGeneration",
+    "OpenAIImageGeneration",
+]