Add FLUX img2img pipeline

Signed-off-by: Daniel Socek <[email protected]>

examples/stable-diffusion/README.md

@@ -610,6 +610,27 @@ python image_to_image_generation.py \
     --bf16
 ```
 
+### FLUX.1 Image to Image
+
+Here is how to generate FLUX.1 images with a single prompt and one input image:
+
+```bash
+python image_to_image_generation.py \
+    --model_name_or_path "black-forest-labs/FLUX.1-dev" \
+    --src_image_path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" \
+    --prompts "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k" \
+    --num_images_per_prompt 20 \
+    --batch_size 4 \
+    --strength 0.9 \
+    --guidance_scale 3.5 \
+    --num_inference_steps 30 \
+    --image_save_dir /tmp/flux_images \
+    --use_habana \
+    --use_hpu_graphs \
+    --gaudi_config Habana/stable-diffusion \
+    --bf16
+```
+
 ### Stable Diffusion Image Variations
 
 Here is how to generate images with one image, it does not accept prompt input

examples/stable-diffusion/image_to_image_generation.py

@@ -27,6 +27,7 @@
     GaudiDDIMScheduler,
     GaudiEulerAncestralDiscreteScheduler,
     GaudiEulerDiscreteScheduler,
+    GaudiFlowMatchEulerDiscreteScheduler,
 )
 from optimum.habana.utils import set_seed
 
@@ -74,7 +75,7 @@ def main():
         type=str,
         nargs="*",
         default=None,
-        help="The second prompt or prompts to guide the image generation (applicable to SDXL).",
+        help="The second prompt or prompts to guide the image generation (applicable to SDXL and FLUX).",
     )
     parser.add_argument(
         "--num_images_per_prompt", type=int, default=1, help="The number of images to generate per prompt."
@@ -101,6 +102,18 @@ def main():
             " of slower inference."
         ),
     )
+    parser.add_argument(
+        "--strength",
+        type=float,
+        default=0.9,
+        help=(
+            "Applicable to FLUX. Indicates extent to transform the reference image. Must be between 0 and 1. Image is used as a"
+            " starting point and more noise is added the higher the `strength`. The number of denoising steps depends"
+            " on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising"
+            " process runs for the full number of iterations specified in `num_inference_steps`. A value of 1 essentially"
+            " ignores reference image."
+        ),
+    )
     parser.add_argument(
         "--guidance_scale",
         type=float,
@@ -210,6 +223,9 @@ def main():
         res["height"] = args.height
     sdxl_models = ["stable-diffusion-xl", "sdxl"]
     sdxl = False
+    flux_models = ["FLUX.1"]
+    flux = False
+
     kwargs = {
         "use_habana": args.use_habana,
         "use_hpu_graphs": args.use_hpu_graphs,
@@ -221,6 +237,10 @@ def main():
         from optimum.habana.diffusers import GaudiStableDiffusionXLImg2ImgPipeline as Img2ImgPipeline
 
         sdxl = True
+    elif any(model in args.model_name_or_path for model in flux_models):
+        from optimum.habana.diffusers import GaudiFluxImg2ImgPipeline as Img2ImgPipeline
+
+        flux = True
     elif "instruct-pix2pix" in args.model_name_or_path:
         from optimum.habana.diffusers import GaudiStableDiffusionInstructPix2PixPipeline as Img2ImgPipeline
 
@@ -274,10 +294,14 @@ def main():
         pipeline.scheduler = GaudiEulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
     elif pipeline.scheduler.config._class_name == "EulerDiscreteScheduler":
         pipeline.scheduler = GaudiEulerDiscreteScheduler.from_config(pipeline.scheduler.config)
+    elif pipeline.scheduler.config._class_name == "FlowMatchEulerDiscreteScheduler":
+        pipeline.scheduler = GaudiFlowMatchEulerDiscreteScheduler.from_config(pipeline.scheduler.config)
     else:
         pipeline.scheduler = GaudiDDIMScheduler.from_config(pipeline.scheduler.config)
+
     # Set seed before running the model
     set_seed(args.seed)
+
     # Generate images
     if sdxl:
         outputs = pipeline(
@@ -296,6 +320,21 @@ def main():
             profiling_steps=args.profiling_steps,
             **res,
         )
+    elif flux:
+        outputs = pipeline(
+            image=image,
+            prompt=args.prompts,
+            prompt_2=args.prompts_2,
+            num_images_per_prompt=args.num_images_per_prompt,
+            batch_size=args.batch_size,
+            num_inference_steps=args.num_inference_steps,
+            strength=args.strength,
+            guidance_scale=args.guidance_scale,
+            output_type=args.output_type,
+            profiling_warmup_steps=args.profiling_warmup_steps,
+            profiling_steps=args.profiling_steps,
+            **res,
+        )
     else:
         outputs = pipeline(
             image=image,

optimum/habana/diffusers/__init__.py

@@ -5,6 +5,7 @@
 )
 from .pipelines.ddpm.pipeline_ddpm import GaudiDDPMPipeline
 from .pipelines.flux.pipeline_flux import GaudiFluxPipeline
+from .pipelines.flux.pipeline_flux_img2img import GaudiFluxImg2ImgPipeline
 from .pipelines.pipeline_utils import GaudiDiffusionPipeline
 from .pipelines.stable_diffusion.pipeline_stable_diffusion import GaudiStableDiffusionPipeline
 from .pipelines.stable_diffusion.pipeline_stable_diffusion_depth2img import GaudiStableDiffusionDepth2ImgPipeline