import base64 import hashlib import os import random from io import BytesIO from typing import Any, Dict, List, Optional from distilabel.llms import InferenceEndpointsLLM # from distilabel.llms.huggingface import InferenceEndpointsLLM from distilabel.pipeline import Pipeline from distilabel.steps import GroupColumns, LoadDataFromHub, StepInput from distilabel.steps.base import StepInput from distilabel.steps.tasks import Task from distilabel.steps.typing import StepOutput from PIL import Image from pydantic import validate_call ## At the time of writing this, the distilabel library does not support the image generation endpoint. ## This is a temporary fix to allow us to use the image generation endpoint. class InferenceEndpointsImageLLM(InferenceEndpointsLLM): @validate_call async def agenerate( self, input: Dict[str, Any], negative_prompt: Optional[str] = None, height: Optional[float] = None, width: Optional[float] = None, num_inference_steps: Optional[float] = None, guidance_scale: Optional[float] = None, ) -> List[Dict[str, Any]]: prompt = input.get("prompt") image = await self._aclient.text_to_image( prompt, negative_prompt=negative_prompt, height=int(height) if height else None, width=int(width) if width else None, num_inference_steps=int(num_inference_steps) if num_inference_steps else None, guidance_scale=float(guidance_scale) if guidance_scale else None, seed=random.randint(0, 1000000), ) buffered = BytesIO() image.save(buffered, format="JPEG") img_str = base64.b64encode(buffered.getvalue()).decode() return [{"image": img_str}] class ImageGeneration(Task): @property def inputs(self) -> List[str]: return ["prompt"] @property def outputs(self) -> List[str]: return ["image", "model_name"] def format_input(self, input: Dict[str, Any]) -> Dict[str, str]: return {"prompt": input["prompt"]} def format_output( self, output: Dict[str, Any], input: Dict[str, Any] ) -> Dict[str, Any]: image_str = output.get("image") image = None if image_str: image_bytes = base64.b64decode(image_str) image = Image.open(BytesIO(image_bytes)) return {"image": image, "model_name": self.llm.model_name} def process(self, *args: StepInput) -> "StepOutput": inputs = args[0] if args else [] formatted_inputs = self._format_inputs(inputs) outputs = self.llm.generate_outputs( inputs=formatted_inputs, num_generations=self.num_generations, **self.llm.get_generation_kwargs(), ) task_outputs = [] for input, input_outputs in zip(inputs, outputs): formatted_outputs = self._format_outputs(input_outputs, input) for formatted_output in formatted_outputs: if "image" in formatted_output and formatted_output["image"]: # use prompt as filename prompt_hash = hashlib.md5(input["prompt"].encode()).hexdigest() self.save_artifact( name="images", write_function=lambda path: formatted_output["image"].save( path / f"{prompt_hash}.jpeg" ), metadata={"type": "image", "library": "diffusers"}, ) formatted_output["image"] = { "path": f"artifacts/{self.name}/images/{prompt_hash}.jpeg" } task_output = { **input, **formatted_output, "model_name": self.llm.model_name, } task_outputs.append(task_output) yield task_outputs ## Let's determine the categories and subcategories for the image generation task # https://huggingface.co/spaces/google/sdxl/blob/main/app.py#L55 categories = { # included "Cinematic": [ # included "emotional", "harmonious", "vignette", "highly detailed", "high budget", "bokeh", "cinemascope", "moody", "epic", "gorgeous", "film grain", "grainy", ], # included "Photographic": [ # included "film", "bokeh", "professional", "4k", "highly detailed", ## not included "Landscape", "Portrait", "Macro", "Portra", "Gold", "ColorPlus", "Ektar", "Superia", "C200", "CineStill", "CineStill 50D", "CineStill 800T", "Tri-X", "HP5", "Delta", "T-Max", "Fomapan", "StreetPan", "Provia", "Ektachrome", "Velvia", ], # included "Anime": [ # included "anime style", "key visual", "vibrant", "studio anime", "highly detailed", ], # included "Manga": [ # included "vibrant", "high-energy", "detailed", "iconic", "Japanese comic style", ], # included "Digital art": [ # included "digital artwork", "illustrative", "painterly", "matte painting", "highly detailed", ], # included "Pixel art": [ # included "low-res", "blocky", "pixel art style", "8-bit graphics", ], # included "Fantasy art": [ # included "magnificent", "celestial", "ethereal", "painterly", "epic", "majestic", "magical", "fantasy art", "cover art", "dreamy", ], # included "Neonpunk": [ # included "cyberpunk", "vaporwave", "neon", "vibes", "vibrant", "stunningly beautiful", "crisp", "detailed", "sleek", "ultramodern", "magenta highlights", "dark purple shadows", "high contrast", "cinematic", "ultra detailed", "intricate", "professional", ], # included "3D Model": [ # included "octane render", "highly detailed", "volumetric", "dramatic lighting", ], # not included "Painting": [ "Oil", "Acrylic", "Watercolor", "Digital", "Mural", "Sketch", "Gouache", "Renaissance", "Baroque", "Romanticism", "Impressionism", "Expressionism", "Cubism", "Surrealism", "Pop Art", "Minimalism", "Realism", "Encaustic", "Tempera", "Fresco", "Ink Wash", "Spray Paint", "Mixed Media", ], # not included "Animation": [ # not included "Animation", "Stop motion", "Claymation", "Pixel Art", "Vector", "Hand-drawn", "Cutout", "Whiteboard", ], # not included "Illustration": [ # not included "Book", "Comics", "Editorial", "Advertising", "Technical", "Fantasy", "Scientific", "Fashion", "Storyboard", "Concept Art", "Manga", "Anime", "Digital", "Vector", "Design", ], } ## We will use the Qwen2.5-72B-Instruct model for the text generation task, this will help us to generate the quality and style prompts sd = InferenceEndpointsImageLLM( base_url="https://el8g78juu06xfxtx.us-east-1.aws.endpoints.huggingface.cloud", api_key=os.getenv("HF_TOKEN"), ) flux_dev = InferenceEndpointsImageLLM( base_url="https://f94i5ss7a040r0v5.us-east-1.aws.endpoints.huggingface.cloud", api_key=os.getenv("HF_TOKEN"), ) ## Let's create the pipeline to generate the quality and style prompts with Pipeline(name="image_preferences_synthetic_data_generation") as pipeline: load_data = LoadDataFromHub(name="load_dataset") image_gen_quality_dev = ImageGeneration( name="image_gen_quality_dev", llm=flux_dev, input_mappings={"prompt": "quality_prompt"}, output_mappings={"image": "image_quality_dev"}, ) image_gen_simplified_dev = ImageGeneration( name="image_gen_simplified_dev", llm=flux_dev, input_mappings={"prompt": "simplified_prompt"}, output_mappings={"image": "image_simplified_dev"}, ) image_gen_quality_sd = ImageGeneration( name="image_gen_quality_sd", llm=sd, input_mappings={"prompt": "quality_prompt"}, output_mappings={"image": "image_quality_sd"}, ) image_gen_simplified_sd = ImageGeneration( name="image_gen_simplified_sd", llm=sd, input_mappings={"prompt": "simplified_prompt"}, output_mappings={"image": "image_simplified_sd"}, ) group_columns_2 = GroupColumns(columns=["model_name"]) ( load_data >> [ image_gen_quality_dev, image_gen_simplified_dev, image_gen_quality_sd, image_gen_simplified_sd, ] >> group_columns_2 ) ## Let's run the pipeline and push the resulting dataset to the hub if __name__ == "__main__": num_examples = 15000 batch_size = 5 num_inference_steps = 25 width = 1024 height = 1024 distiset = pipeline.run( use_cache=True, parameters={ load_data.name: { "num_examples": num_examples, "repo_id": "data-is-better-together/imgsys-results-prompts-style_v2_part1_cleaned", }, image_gen_quality_sd.name: { "llm": { "generation_kwargs": { "width": width, "height": height, "guidance_scale": 4.5, "num_inference_steps": num_inference_steps, }, }, "input_batch_size": batch_size, }, image_gen_quality_dev.name: { "llm": { "generation_kwargs": { "guidance_scale": 4.5, "num_inference_steps": num_inference_steps, }, }, "input_batch_size": batch_size, }, image_gen_simplified_sd.name: { "llm": { "generation_kwargs": { "width": width, "height": height, "guidance_scale": 3.5, "num_inference_steps": num_inference_steps, }, }, "input_batch_size": batch_size, }, image_gen_simplified_dev.name: { "llm": { "generation_kwargs": { "guidance_scale": 3.5, "num_inference_steps": num_inference_steps, }, }, "input_batch_size": batch_size, }, }, ) from pathlib import Path from PIL import Image dataset = distiset["default"]["train"] artifacts_path = Path(distiset.artifacts_path) def load_images(batch): for column in [ "image_quality_dev", "image_simplified_dev", "image_quality_sd", "image_simplified_sd", ]: batch[f"{column}_loaded"] = [] for i in range(len(batch[column])): if batch[column][i]: batch[f"{column}_loaded"].append( Image.open( artifacts_path / batch[column][i]["path"].replace("artifacts/", "") ) ) else: batch[f"{column}_loaded"].append(None) return batch dataset = dataset.map(load_images, batched=True) dataset = dataset.remove_columns( [ "image_quality_dev", "image_simplified_dev", "image_quality_sd", "image_simplified_sd", ] ) dataset = dataset.rename_columns( { "image_quality_dev_loaded": "image_quality_dev", "image_simplified_dev_loaded": "image_simplified_dev", "image_quality_sd_loaded": "image_quality_sd", "image_simplified_sd_loaded": "image_simplified_sd", } ) distiset["default"]["train"] = dataset distiset.artifacts_path = "" distiset.push_to_hub( repo_id="data-is-better-together/image-preferences", include_script=True, generate_card=False, token=os.getenv("HF_TOKEN"), )