data_juicer.ops.mapper.sdxl_prompt2prompt_mapper module¶
- class data_juicer.ops.mapper.sdxl_prompt2prompt_mapper.SDXLPrompt2PromptMapper(hf_diffusion: str = 'stabilityai/stable-diffusion-xl-base-1.0', trust_remote_code=False, torch_dtype: str = 'fp32', num_inference_steps: float = 50, guidance_scale: float = 7.5, text_key=None, text_key_second=None, output_dir='/home/runner/.cache/data_juicer/assets', *args, **kwargs)[source]¶
Bases:
Mapper
Generates pairs of similar images using the SDXL model.
This operator uses a Hugging Face diffusion model to generate image pairs based on two text prompts. The quality and similarity of the generated images are controlled by parameters such as num_inference_steps and guidance_scale. The first and second text prompts are specified using text_key and text_key_second, respectively. The generated images are saved in the specified output_dir with unique filenames. The operator requires both text keys to be set for processing.
- __init__(hf_diffusion: str = 'stabilityai/stable-diffusion-xl-base-1.0', trust_remote_code=False, torch_dtype: str = 'fp32', num_inference_steps: float = 50, guidance_scale: float = 7.5, text_key=None, text_key_second=None, output_dir='/home/runner/.cache/data_juicer/assets', *args, **kwargs)[source]¶
Initialization method.
- Parameters:
hf_diffusion – diffusion model name on huggingface to generate the image.
trust_remote_code – whether to trust the remote code of HF models.
torch_dtype – the floating point type used to load the diffusion model.
num_inference_steps – The larger the value, the better the
image generation quality; however, this also increases the time required for generation. :param guidance_scale: A higher guidance scale value encourages the
model to generate images closely linked to the text prompt at the expense of lower image quality. Guidance scale is enabled when
- Parameters:
text_key – the key name used to store the first caption in the caption pair.
text_key_second – the key name used to store the second caption in the caption pair.
output_dir – the storage location of the generated images.