diff --git a/examples/stable-diffusion/README.md b/examples/stable-diffusion/README.md
index b922b92498..fb4bdee33d 100644
--- a/examples/stable-diffusion/README.md
+++ b/examples/stable-diffusion/README.md
@@ -28,12 +28,12 @@ First, you should install the requirements:
 pip install -r requirements.txt
 ```
 
-
 ## Text-to-image Generation
 
 ### Single Prompt
 
 Here is how to generate images with one prompt:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path CompVis/stable-diffusion-v1-4 \
@@ -51,10 +51,10 @@ python text_to_image_generation.py \
 > The first batch of images entails a performance penalty. All subsequent batches will be generated much faster.
 > You can enable this mode with `--use_hpu_graphs`.
 
-
 ### Multiple Prompts
 
 Here is how to generate images with several prompts:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path CompVis/stable-diffusion-v1-4 \
@@ -69,7 +69,9 @@ python text_to_image_generation.py \
 ```
 
 ### Distributed inference with multiple HPUs
+
 Here is how to generate images with two prompts on two HPUs:
+
 ```bash
 python ../gaudi_spawn.py \
     --world_size 2 text_to_image_generation.py \
@@ -109,10 +111,10 @@ python text_to_image_generation.py \
 ```
 
 > There are two different checkpoints for Stable Diffusion 2:
+>
 > - use [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) for generating 768x768 images
 > - use [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base) for generating 512x512 images
 
-
 ### Latent Diffusion Model for 3D (LDM3D)
 
 [LDM3D](https://arxiv.org/abs/2305.10853) generates both image and depth map data from a given text prompt, allowing users to generate RGBD images from text prompts.
@@ -135,7 +137,9 @@ python text_to_image_generation.py \
     --ldm3d \
     --bf16
 ```
+
 Here is how to generate images and depth maps with two prompts on two HPUs:
+
 ```bash
 python ../gaudi_spawn.py \
     --world_size 2 text_to_image_generation.py \
@@ -154,6 +158,7 @@ python ../gaudi_spawn.py \
 ```
 
 > There are three different checkpoints for LDM3D:
+>
 > - use [original checkpoint](https://huggingface.co/Intel/ldm3d) to generate outputs from the paper
 > - use [the latest checkpoint](https://huggingface.co/Intel/ldm3d-4c) for generating improved results
 > - use [the pano checkpoint](https://huggingface.co/Intel/ldm3d-pano) to generate panoramic view
@@ -163,6 +168,7 @@ python ../gaudi_spawn.py \
 Stable Diffusion XL was proposed in [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://arxiv.org/pdf/2307.01952.pdf) by the Stability AI team.
 
 Here is how to generate SDXL images with a single prompt:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \
@@ -182,6 +188,7 @@ python text_to_image_generation.py \
 > You can enable this mode with `--use_hpu_graphs`.
 
 Here is how to generate SDXL images with several prompts:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \
@@ -199,6 +206,7 @@ python text_to_image_generation.py \
 SDXL combines a second text encoder (OpenCLIP ViT-bigG/14) with the original text encoder to significantly
 increase the number of parameters. Here is how to generate images with several prompts for both `prompt`
 and `prompt_2` (2nd text encoder), as well as their negative prompts:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \
@@ -217,6 +225,7 @@ python text_to_image_generation.py \
 ```
 
 Here is how to generate SDXL images with two prompts on two HPUs:
+
 ```bash
 python ../gaudi_spawn.py \
     --world_size 2 text_to_image_generation.py \
@@ -235,14 +244,17 @@ python ../gaudi_spawn.py \
     --bf16 \
     --distributed
 ```
+
 > HPU graphs are recommended when generating images by batches to get the fastest possible generations.
 > The first batch of images entails a performance penalty. All subsequent batches will be generated much faster.
 > You can enable this mode with `--use_hpu_graphs`.
 
 ### SDXL-Turbo
+
 SDXL-Turbo is a distilled version of SDXL 1.0, trained for real-time synthesis.
 
 Here is how to generate images with multiple prompts:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path stabilityai/sdxl-turbo \
@@ -275,11 +287,13 @@ Before running SD3 pipeline, you need to:
 
 1. Agree to the Terms and Conditions for using SD3 model at [HuggingFace model page](https://huggingface.co/stabilityai/stable-diffusion-3-medium)
 2. Authenticate with HuggingFace using your HF Token. For authentication, run:
+
 ```bash
 huggingface-cli login
 ```
 
 Here is how to generate SD3 images with a single prompt:
+
 ```bash
 PT_HPU_MAX_COMPOUND_OP_SIZE=1 \
 python text_to_image_generation.py \
@@ -299,12 +313,32 @@ python text_to_image_generation.py \
 > For improved performance of the SD3 pipeline on Gaudi, it is recommended to configure the environment
 > by setting PT_HPU_MAX_COMPOUND_OP_SIZE to 1.
 
+### FLUX.1
+
+FLUX.1 was was introduced by Black Forest Labs [here](https://blackforestlabs.ai/announcing-black-forest-labs/)
+
+```bash
+python text_to_image_generation.py \
+    --model_name_or_path black-forest-labs/FLUX.1-schnell \
+    --prompts "A cat holding a sign that says hello world" \
+    --num_images_per_prompt 10 \
+    --batch_size 1 \
+    --num_inference_steps 28 \
+    --image_save_dir /tmp/flux_1_images \
+    --scheduler flow_match_euler_discrete\
+    --use_habana \
+    --use_hpu_graphs \
+    --gaudi_config Habana/stable-diffusion \
+    --bf16
+```
+
 ## ControlNet
 
-ControlNet was introduced in [Adding Conditional Control to Text-to-Image Diffusion Models ](https://huggingface.co/papers/2302.05543) by Lvmin Zhang and Maneesh Agrawala.
+ControlNet was introduced in [Adding Conditional Control to Text-to-Image Diffusion Models](https://huggingface.co/papers/2302.05543) by Lvmin Zhang and Maneesh Agrawala.
 It is a type of model for controlling StableDiffusion by conditioning the model with an additional input image.
 
 Here is how to generate images conditioned by canny edge model:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path CompVis/stable-diffusion-v1-4 \
@@ -321,6 +355,7 @@ python text_to_image_generation.py \
 ```
 
 Here is how to generate images conditioned by canny edge model and with multiple prompts:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path CompVis/stable-diffusion-v1-4 \
@@ -337,6 +372,7 @@ python text_to_image_generation.py \
 ```
 
 Here is how to generate images conditioned by canny edge model and with two prompts on two HPUs:
+
 ```bash
 python ../gaudi_spawn.py \
     --world_size 2 text_to_image_generation.py \
@@ -355,6 +391,7 @@ python ../gaudi_spawn.py \
 ```
 
 Here is how to generate images conditioned by open pose model:
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path CompVis/stable-diffusion-v1-4 \
@@ -372,6 +409,7 @@ python text_to_image_generation.py \
 ```
 
 Here is how to generate images with conditioned by canny edge model using Stable Diffusion 2
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path stabilityai/stable-diffusion-2-1 \
@@ -395,6 +433,7 @@ Inpainting replaces or edits specific areas of an image. For more details,
 please refer to [Hugging Face Diffusers doc](https://huggingface.co/docs/diffusers/en/using-diffusers/inpaint).
 
 ### Stable Diffusion Inpainting
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path  stabilityai/stable-diffusion-2-inpainting \
@@ -412,6 +451,7 @@ python text_to_image_generation.py \
 ```
 
 ### Stable Diffusion XL Inpainting
+
 ```bash
 python text_to_image_generation.py \
     --model_name_or_path  diffusers/stable-diffusion-xl-1.0-inpainting-0.1\
@@ -457,10 +497,10 @@ python image_to_image_generation.py \
 > The first batch of images entails a performance penalty. All subsequent batches will be generated much faster.
 > You can enable this mode with `--use_hpu_graphs`.
 
-
 ### Multiple Prompts
 
 Here is how to generate images with several prompts and one image.
+
 ```bash
 python image_to_image_generation.py \
     --model_name_or_path "timbrooks/instruct-pix2pix" \
@@ -482,10 +522,10 @@ python image_to_image_generation.py \
 > The first batch of images entails a performance penalty. All subsequent batches will be generated much faster.
 > You can enable this mode with `--use_hpu_graphs`.
 
-
 ### Stable Diffusion XL Refiner
 
 Here is how to generate SDXL images with a single prompt and one image:
+
 ```bash
 python image_to_image_generation.py \
     --model_name_or_path "stabilityai/stable-diffusion-xl-refiner-1.0" \
@@ -505,6 +545,7 @@ python image_to_image_generation.py \
 ### Stable Diffusion Image Variations
 
 Here is how to generate images with one image, it does not accept prompt input
+
 ```bash
 python image_to_image_generation.py \
     --model_name_or_path "lambdalabs/sd-image-variations-diffusers" \
@@ -625,6 +666,7 @@ Script `image_to_video_generation.py` showcases how to perform image-to-video ge
 ### Single Image Prompt
 
 Here is how to generate video with one image prompt:
+
 ```bash
 PT_HPU_MAX_COMPOUND_OP_SIZE=1 \
 python image_to_video_generation.py \
@@ -645,6 +687,7 @@ python image_to_video_generation.py \
 ### Multiple Image Prompts
 
 Here is how to generate videos with several image prompts:
+
 ```bash
 PT_HPU_MAX_COMPOUND_OP_SIZE=1 \
 python image_to_video_generation.py \
diff --git a/examples/stable-diffusion/prompts_100.txt b/examples/stable-diffusion/prompts_100.txt
new file mode 100644
index 0000000000..09a94f1b9a
--- /dev/null
+++ b/examples/stable-diffusion/prompts_100.txt
@@ -0,0 +1,100 @@
+A women playing tennis on a blue tennis court.
+Two surfers in wetsuits carrying surfboards along the beach.
+People are flying their kites in a large field.
+A statue that is in front of a building.
+A man attempting to do a skateboard trick on an outdoor halfpipe.
+A cream bathroom with red accents and open window.
+A woman in a white dress and a man in gray stand near a cake on a white table under a white canopy.
+A tour bus downtown with yoga ads all over it.
+Three people riding horses on a beach next to the ocean..
+A bear lying in its den on a pile of wood.
+A herd of cattle is feeding at the river's edge.
+there is a male snowboarder that is in the air
+A statue of an elephant with tattoos and a target drawing
+a woman is standing over a white cake
+A grey motorcycle parked in a tropical setting.
+Two dogs playing in the grass with a frisbee.
+A group of people standing around a table full of food.
+A person holding a piece of broccoli with an insect on it. 
+A man in a gray suit and a red tie.
+Someone with skis on his back walking up a snow covered mountain.
+a woman in a blue shirt holding a pair of large scissors
+A fluffy white cat has a frowning look on it's face.
+A bus that is sitting on the street.
+Sinks in the washroom that is public and white.
+this lady is walking along the shore on a beach
+A woman helping a man to do his tie.
+A tray with coffee and a pastry on it.
+A man with a bald head and a bear wearing a bow tie.
+A cat eating a birthday cake on top of the table. 
+A man at a party talking on a cell phone.
+Three commuter buses sitting outside of a building.
+Food truck with customers ordering them with friends.
+A couple of bikes in front of a small stone wall.
+Someone is enjoying a small slice of pie. 
+there is a small tv and coffee table in the living room
+A plate with two sandwiches, cup and knife on the table
+A group of people venturing out on a horseback ride.
+A horse walking through a grassy field while two cows eat hay. 
+A bathroom scene with focus on the toilet.
+A man on a field swinging a baseball bat.
+Some people are standing on a crowd crowded sidewalk
+A pregnant women taking a picture of herself in the mirror.
+A line of police offices riding horses down a street.
+A buffet styled restaurant without self service but a server.
+A tview of a living room with fold out bed.
+A surfboard advertising offerings as people check them out.
+a bath room with a sink a mirror and towel racks
+A desk with a computer on in and a key board
+A large clock mounted on the wall of a stone building
+Three men sitting around a table with wine on it. 
+Man with glasses and a mustache standing in front of a door.
+A woman in a black dress holding a racquet.
+This hotel room has a king size bed. 
+A severely injured man hooked up to machines  in the hospital
+A black and white picture of an old store.
+Fingers keep a meatball sub from falling apart.
+A plate of fish covered in marinara, cheese, carrots, a fork, next to bread.
+A anal filled with boats and the street above it filled with people under umbrellas. 
+A woman standing on a tennis court holding a racquet.
+An elephant is standing next to a tree and a fence.
+People in a street with birds all over.
+THERE ARE PEOPLE THAT ARE STANDING IN THE GRASS
+a man standing by a fence while throwing a frisbee
+A pot that is on the stove with some food in it.
+A stuffed animal is inside of a microwave.
+Group of parents watching small children on a baseball field. 
+The kite is flying  high in the air 
+This cat is playing on a fuzzy white blanket.
+A lone giraffe at a zoo with trees behind it.
+A train engine carrying carts down a track.
+a young man brushes his teeth in the bathroom
+Two children playing baseball in red uniforms and hats.
+Men lined up an a runway in a desert greet an arriving jet plane.
+The American flag flies next to the clock tower on a snowy day. 
+A man takes a selfie of himself in the mirror. 
+A baseball game where a player is running to 3rd base.
+A man and a woman standing in front of a bus.
+A plate of dessert sitting beside a drink in a cafe.
+A man on the couch is petting the dog
+A cat is sleeping with a remote control on a couch.
+A European fighter jet flying above the tree tops.
+Snowman's head has a carrot for a nose and lemon slices for eyes.
+A man standing on a tennis court holding a tennis racquet.
+Two large green and white jumbo jet planes on the tarmac.
+The bedroom is is decorated in various zebra prints.
+Dog displaying skills near disc in open grassy area.
+A toilet facility in a stone cell on a plank floor.
+A picture done by Independent Expression Photography of a girl posing in an empty road sitting on her suit cases.
+Pair of colorful stuffed bears hanging on line in backyard.
+a large train is on the track going by the ocean
+There is a pizza with olives, peppers, meat, and cheese on the table.
+A man standing in front of microphones. 
+A woman that is holding a book sitting on a bed.
+A man showing a ring at a formal event.
+A concrete building with towers, a steep in the middle and a clock underneath.
+A white dog standing on top of a wooden bench.
+A woman holds a plate with rainbow cake.
+A women who is taking a picture of her food.
+A man with long hair and in a towel holding a toothbrush.
+a woman holding a tennis racket in the air 
diff --git a/examples/stable-diffusion/prompts_5.txt b/examples/stable-diffusion/prompts_5.txt
new file mode 100644
index 0000000000..6254d2714e
--- /dev/null
+++ b/examples/stable-diffusion/prompts_5.txt
@@ -0,0 +1,5 @@
+A women playing tennis on a blue tennis court.
+Two surfers in wetsuits carrying surfboards along the beach.
+People are flying their kites in a large field.
+A statue that is in front of a building.
+A man attempting to do a skateboard trick on an outdoor halfpipe.
diff --git a/examples/stable-diffusion/quantize/measure_config.json b/examples/stable-diffusion/quantize/measure_config.json
new file mode 100755
index 0000000000..19d5e988d4
--- /dev/null
+++ b/examples/stable-diffusion/quantize/measure_config.json
@@ -0,0 +1,5 @@
+{
+    "method": "HOOKS",
+    "mode": "MEASURE",
+    "dump_stats_path": "quantize/measure_all/fp8"
+}
diff --git a/examples/stable-diffusion/quantize/quant_config.json b/examples/stable-diffusion/quantize/quant_config.json
new file mode 100755
index 0000000000..eab3011a5e
--- /dev/null
+++ b/examples/stable-diffusion/quantize/quant_config.json
@@ -0,0 +1,6 @@
+{
+    "method": "HOOKS",
+    "mode": "QUANTIZE",
+    "scale_method": "maxabs_hw_opt_weight",
+    "dump_stats_path": "quantize/measure_all/fp8"
+}
diff --git a/examples/stable-diffusion/quantize/quant_config_500.json b/examples/stable-diffusion/quantize/quant_config_500.json
new file mode 100755
index 0000000000..173f93772c
--- /dev/null
+++ b/examples/stable-diffusion/quantize/quant_config_500.json
@@ -0,0 +1,6 @@
+{
+    "method": "HOOKS",
+    "mode": "QUANTIZE",
+    "scale_method": "maxabs_hw_opt_weight",
+    "dump_stats_path": "quantize/measure_all_500/fp8"
+}
diff --git a/examples/stable-diffusion/quantize/quant_config_bmm.json b/examples/stable-diffusion/quantize/quant_config_bmm.json
new file mode 100755
index 0000000000..5cbc2bac9f
--- /dev/null
+++ b/examples/stable-diffusion/quantize/quant_config_bmm.json
@@ -0,0 +1,7 @@
+{
+    "method": "HOOKS",
+    "mode": "QUANTIZE",
+    "scale_method": "maxabs_hw_opt_weight",
+    "dump_stats_path": "quantize/measure_all/fp8",
+    "blocklist": {"types": ["Linear", "Conv2d", "LoRACompatibleLinear", "LoRACompatibleConv"]}
+}
diff --git a/examples/stable-diffusion/text_to_image_generation.py b/examples/stable-diffusion/text_to_image_generation.py
index 0d8ff554d1..fed390fc12 100755
--- a/examples/stable-diffusion/text_to_image_generation.py
+++ b/examples/stable-diffusion/text_to_image_generation.py
@@ -26,7 +26,7 @@
 from optimum.habana.diffusers import (
     GaudiDDIMScheduler,
     GaudiEulerAncestralDiscreteScheduler,
-    GaudiEulerDiscreteScheduler,
+    GaudiEulerDiscreteScheduler
 )
 from optimum.habana.utils import set_seed
 
@@ -66,7 +66,7 @@ def main():
     parser.add_argument(
         "--scheduler",
         default="ddim",
-        choices=["default", "euler_discrete", "euler_ancestral_discrete", "ddim"],
+        choices=["default", "euler_discrete", "euler_ancestral_discrete", "ddim", "flow_match_euler_discrete"],
         type=str,
         help="Name of scheduler",
     )
@@ -286,13 +286,27 @@ def main():
         action="store_true",
         help="Use rescale_betas_zero_snr for controlling image brightness",
     )
+    parser.add_argument(
+        "--quant_mode",
+        default="disable",
+        type=str,
+        help="Quantization mode 'measure', 'quantize', 'quantize-mixed' or 'disable'",
+    )
+    parser.add_argument(
+        "--prompts_file",
+        type=str,
+        default=None,
+        help="The file with prompts (for large number of images generation).",
+    )
     args = parser.parse_args()
 
     # Select stable diffuson pipeline based on input
     sdxl_models = ["stable-diffusion-xl", "sdxl"]
     sd3_models = ["stable-diffusion-3"]
+    flux_models = ["FLUX.1-dev", "FLUX.1-schnell","OpenFLUX.1"]
     sdxl = True if any(model in args.model_name_or_path for model in sdxl_models) else False
     sd3 = True if any(model in args.model_name_or_path for model in sd3_models) else False
+    flux = True if any(model in args.model_name_or_path for model in flux_models) else False
     controlnet = True if args.control_image is not None else False
     inpainting = True if (args.base_image is not None) and (args.mask_image is not None) else False
 
@@ -356,16 +370,18 @@ def main():
                 negative_prompts = negative_prompt
     kwargs_call["negative_prompt"] = negative_prompts
 
-    if sdxl or sd3:
+    if sdxl or sd3 or flux:
         prompts_2 = args.prompts_2
-        negative_prompts_2 = args.negative_prompts_2
         if args.distributed and args.prompts_2 is not None:
             with distributed_state.split_between_processes(args.prompts_2) as prompt_2:
                 prompts_2 = prompt_2
+        kwargs_call["prompt_2"] = prompts_2
+
+    if sdxl or sd3:
+        negative_prompts_2 = args.negative_prompts_2
         if args.distributed and args.negative_prompts_2 is not None:
             with distributed_state.split_between_processes(args.negative_prompts_2) as negative_prompt_2:
                 negative_prompts_2 = negative_prompt_2
-        kwargs_call["prompt_2"] = prompts_2
         kwargs_call["negative_prompt_2"] = negative_prompts_2
 
     if sd3:
@@ -404,7 +420,10 @@ def main():
             control_image = Image.fromarray(image)
         kwargs_call["image"] = control_image
 
+    kwargs_call["quant_mode"] = args.quant_mode
+
     # Instantiate a Stable Diffusion pipeline class
+    import habana_frameworks.torch.core as htcore
     if sdxl:
         # SDXL pipelines
         if controlnet:
@@ -446,6 +465,22 @@ def main():
                 args.model_name_or_path,
                 **kwargs,
             )
+    elif flux:
+        # Flux pipelines
+        if controlnet:
+            # Import Flux+ControlNet pipeline
+            raise ValueError("Flux+ControlNet pipeline is not currenly supported")
+        elif inpainting:
+            # Import FLux Inpainting pipeline
+            raise ValueError("Flux Inpainting pipeline is not currenly supported")
+        else:
+            # Import Flux pipeline
+            from optimum.habana.diffusers import GaudiFluxPipeline
+
+            pipeline = GaudiFluxPipeline.from_pretrained(
+                args.model_name_or_path,
+                **kwargs,
+            )
 
     else:
         # SD pipelines (SD1.x, SD2.x)
@@ -538,6 +573,14 @@ def main():
 
         pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.5, b2=1.6)
 
+    # If prompts file is specified override prompts from the file
+    if args.prompts_file is not None:
+        lines = []
+        with open(args.prompts_file, "r") as file:
+            lines = file.readlines()
+        lines = [line.strip() for line in lines]
+        args.prompts = lines
+
     # Generate Images using a Stable Diffusion pipeline
     if args.distributed:
         with distributed_state.split_between_processes(args.prompts) as prompt:
diff --git a/examples/stable-diffusion/unconditional_image_generation.py b/examples/stable-diffusion/unconditional_image_generation.py
old mode 100644
new mode 100755
diff --git a/optimum/habana/diffusers/__init__.py b/optimum/habana/diffusers/__init__.py
index de76c24f5e..a87833fcca 100644
--- a/optimum/habana/diffusers/__init__.py
+++ b/optimum/habana/diffusers/__init__.py
@@ -4,6 +4,7 @@
     GaudiStableVideoDiffusionControlNetPipeline,
 )
 from .pipelines.ddpm.pipeline_ddpm import GaudiDDPMPipeline
+from .pipelines.flux.pipeline_flux import GaudiFluxPipeline
 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
@@ -23,4 +24,4 @@
 from .pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint import GaudiStableDiffusionXLInpaintPipeline
 from .pipelines.stable_video_diffusion.pipeline_stable_video_diffusion import GaudiStableVideoDiffusionPipeline
 from .pipelines.text_to_video_synthesis.pipeline_text_to_video_synth import GaudiTextToVideoSDPipeline
-from .schedulers import GaudiDDIMScheduler, GaudiEulerAncestralDiscreteScheduler, GaudiEulerDiscreteScheduler
+from .schedulers import GaudiDDIMScheduler, GaudiEulerAncestralDiscreteScheduler, GaudiEulerDiscreteScheduler, GaudiFlowMatchEulerDiscreteScheduler
diff --git a/optimum/habana/diffusers/models/attention_processor.py b/optimum/habana/diffusers/models/attention_processor.py
index b0461a272b..2202f517ec 100755
--- a/optimum/habana/diffusers/models/attention_processor.py
+++ b/optimum/habana/diffusers/models/attention_processor.py
@@ -19,6 +19,7 @@
 import torch
 import torch.nn.functional as F
 from diffusers.models.attention_processor import Attention
+from diffusers.models.embeddings import apply_rotary_emb
 from diffusers.utils import USE_PEFT_BACKEND, logging
 from diffusers.utils.import_utils import is_xformers_available
 from torch import nn
@@ -186,4 +187,191 @@ def __call__(
         return hidden_states
 
 
-AttentionProcessor = Union[AttnProcessor2_0,]
+class GaudiFluxAttnProcessor2_0:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        if encoder_hidden_states is not None:
+            # `context` projections.
+            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        # hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+        from habana_frameworks.torch.hpex.kernels import FusedSDPA
+        import habana_frameworks.torch.hpu as ht
+        hidden_states = FusedSDPA.apply(query, key, value, None, 0.0, False, None, 'fast', None)
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1]:],
+            )
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+
+class GaudiFusedFluxAttnProcessor2_0:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "FusedFluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        qkv = attn.to_qkv(hidden_states)
+        split_size = qkv.shape[-1] // 3
+        query, key, value = torch.split(qkv, split_size, dim=-1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        # `context` projections.
+        if encoder_hidden_states is not None:
+            encoder_qkv = attn.to_added_qkv(encoder_hidden_states)
+            split_size = encoder_qkv.shape[-1] // 3
+            (
+                encoder_hidden_states_query_proj,
+                encoder_hidden_states_key_proj,
+                encoder_hidden_states_value_proj,
+            ) = torch.split(encoder_qkv, split_size, dim=-1)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        # hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+        from habana_frameworks.torch.hpex.kernels import FusedSDPA
+        import habana_frameworks.torch.hpu as ht
+        hidden_states = FusedSDPA.apply(query, key, value, None, 0.0, False, None, 'fast', None)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1]:],
+            )
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+
+AttentionProcessor = Union[AttnProcessor2_0, GaudiFluxAttnProcessor2_0, GaudiFusedFluxAttnProcessor2_0]
diff --git a/optimum/habana/diffusers/pipelines/auto_pipeline.py b/optimum/habana/diffusers/pipelines/auto_pipeline.py
index 77171c9502..a7fb5431ca 100644
--- a/optimum/habana/diffusers/pipelines/auto_pipeline.py
+++ b/optimum/habana/diffusers/pipelines/auto_pipeline.py
@@ -33,6 +33,8 @@
 from .stable_diffusion.pipeline_stable_diffusion_inpaint import GaudiStableDiffusionInpaintPipeline
 from .stable_diffusion_xl.pipeline_stable_diffusion_xl import GaudiStableDiffusionXLPipeline
 from .stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint import GaudiStableDiffusionXLInpaintPipeline
+from .stable_diffusion_3.pipeline_stable_diffusion_3 import GaudiStableDiffusion3Pipeline
+from .flux.pipeline_flux import GaudiFluxPipeline
 
 
 GAUDI_PREFIX_NAME = "Gaudi"
@@ -42,6 +44,8 @@
         ("stable-diffusion", GaudiStableDiffusionPipeline),
         ("stable-diffusion-xl", GaudiStableDiffusionXLPipeline),
         ("stable-diffusion-controlnet", GaudiStableDiffusionControlNetPipeline),
+        ("stable-diffusion-3", GaudiStableDiffusion3Pipeline),
+        ("flux", GaudiFluxPipeline),
     ]
 )
 
diff --git a/optimum/habana/diffusers/pipelines/flux/pipeline_flux.py b/optimum/habana/diffusers/pipelines/flux/pipeline_flux.py
new file mode 100644
index 0000000000..d0d9fd7c35
--- /dev/null
+++ b/optimum/habana/diffusers/pipelines/flux/pipeline_flux.py
@@ -0,0 +1,640 @@
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# 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 diffusers.models.attention_processor import Attention
+import torch.nn.functional as F
+import math
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+from diffusers.utils import BaseOutput, replace_example_docstring
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.models.transformers import FluxTransformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.pipelines.flux.pipeline_flux import FluxPipeline, calculate_shift, retrieve_timesteps
+
+from optimum.utils import logging
+
+from ....transformers.gaudi_configuration import GaudiConfig
+from ....utils import HabanaProfile, speed_metrics, warmup_inference_steps_time_adjustment
+from ..pipeline_utils import GaudiDiffusionPipeline
+from ...models.attention_processor import GaudiFluxAttnProcessor2_0, GaudiFusedFluxAttnProcessor2_0
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class GaudiFluxPipelineOutput(BaseOutput):
+    """
+    Output class for Stable Diffusion pipelines.
+
+    Args:
+        images (`List[PIL.Image.Image]` or `np.ndarray`)
+            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]
+    throughput: float
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from optimum.habana.diffusers import GaudiFluxPipeline
+
+        >>> pipe = GaudiFluxPipeline.from_pretrained(
+        ...    "black-forest-labs/FLUX.1-schnell",
+        ...     torch_dtype=torch.bfloat16,
+        ...     use_habana=True,
+        ...     use_hpu_graphs=True,
+        ...     gaudi_config="Habana/stable-diffusion",
+        ... )
+        >>> prompt = "A cat holding a sign that says hello world"
+        >>> # Depending on the variant being used, the pipeline call will slightly vary.
+        >>> # Refer to the pipeline documentation for more details.
+        >>> image = pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0]
+        >>> image.save("flux.png")
+        ```
+"""
+
+
+class GaudiFluxPipeline(GaudiDiffusionPipeline, FluxPipeline):
+    r"""
+    Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux.py#L140
+
+    The Flux pipeline for text-to-image generation.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Args:
+        transformer ([`FluxTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_2: T5TokenizerFast,
+        transformer: FluxTransformer2DModel,
+        use_habana: bool = False,
+        use_hpu_graphs: bool = False,
+        gaudi_config: Union[str, GaudiConfig] = None,
+        bf16_full_eval: bool = False,
+    ):
+        GaudiDiffusionPipeline.__init__(
+            self,
+            use_habana,
+            use_hpu_graphs,
+            gaudi_config,
+            bf16_full_eval,
+        )
+        FluxPipeline.__init__(
+            self,
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            text_encoder_2=text_encoder_2,
+            tokenizer_2=tokenizer_2,
+            transformer=transformer,
+        )
+
+        for block in self.transformer.single_transformer_blocks:
+            block.attn.processor = GaudiFluxAttnProcessor2_0()
+        for block in self.transformer.transformer_blocks:
+            block.attn.processor = GaudiFluxAttnProcessor2_0()
+
+        self.to(self._device)
+        if use_hpu_graphs:
+            from habana_frameworks.torch.hpu import wrap_in_hpu_graph
+
+            transformer = wrap_in_hpu_graph(transformer)
+
+    @classmethod
+    def _split_inputs_into_batches(cls, batch_size, latents, prompt_embeds, pooled_prompt_embeds, text_ids, latent_image_ids, guidance):
+        # Use torch.split to generate num_batches batches of size batch_size
+        latents_batches = list(torch.split(latents, batch_size))
+        prompt_embeds_batches = list(torch.split(prompt_embeds, batch_size))
+        if pooled_prompt_embeds is not None:
+            pooled_prompt_embeds_batches = list(torch.split(pooled_prompt_embeds, batch_size))
+        if text_ids is not None and text_ids.ndim == 3:
+            text_ids_batches = list(torch.split(text_ids, batch_size))
+        if latent_image_ids is not None and latent_image_ids.ndim == 3:
+            latent_image_ids_batches = list(torch.split(latent_image_ids, batch_size))
+        if guidance is not None:
+            guidance_batches = list(torch.split(guidance, batch_size))
+        else:
+            guidance_batches = [torch.tensor(float('nan')),] * len(latents_batches)
+
+        # If the last batch has less samples than batch_size, pad it with dummy samples
+        num_dummy_samples = 0
+        if latents_batches[-1].shape[0] < batch_size:
+            num_dummy_samples = batch_size - latents_batches[-1].shape[0]
+
+            # Pad latents_batches
+            sequence_to_stack = (latents_batches[-1],) + tuple(
+                torch.zeros_like(latents_batches[-1][0][None, :]) for _ in range(num_dummy_samples)
+            )
+            latents_batches[-1] = torch.vstack(sequence_to_stack)
+
+            # Pad prompt_embeds_batches
+            sequence_to_stack = (prompt_embeds_batches[-1],) + tuple(
+                torch.zeros_like(prompt_embeds_batches[-1][0][None, :]) for _ in range(num_dummy_samples)
+            )
+            prompt_embeds_batches[-1] = torch.vstack(sequence_to_stack)
+
+            # Pad pooled_prompt_embeds if necessary
+            if pooled_prompt_embeds is not None:
+                sequence_to_stack = (pooled_prompt_embeds_batches[-1],) + tuple(
+                    torch.zeros_like(pooled_prompt_embeds_batches[-1][0][None, :]) for _ in range(num_dummy_samples)
+                )
+                pooled_prompt_embeds_batches[-1] = torch.vstack(sequence_to_stack)
+
+            # Pad text_ids_batches if necessary
+            if text_ids is not None and text_ids.ndim == 3:
+                sequence_to_stack = (text_ids_batches[-1],) + tuple(
+                    torch.zeros_like(text_ids_batches[-1][0][None, :]) for _ in range(num_dummy_samples)
+                )
+                text_ids_batches[-1] = torch.vstack(sequence_to_stack)
+
+            # Pad latent_image_ids if necessary
+            if latent_image_ids is not None and latent_image_ids.ndim == 3:
+                sequence_to_stack = (latent_image_ids_batches[-1],) + tuple(
+                    torch.zeros_like(latent_image_ids_batches[-1][0][None, :]) for _ in range(num_dummy_samples)
+                )
+                latent_image_ids_batches[-1] = torch.vstack(sequence_to_stack)
+
+            # Pad guidance if necessary
+            if guidance is not None:
+                sequence_to_stack = (guidance_batches[-1],) + tuple(
+                    torch.zeros_like(guidance_batches[-1][0][None, :]) for _ in range(num_dummy_samples)
+                )
+                guidance_batches[-1] = torch.vstack(sequence_to_stack)
+
+        # Stack batches in the same tensor
+        latents_batches = torch.stack(latents_batches)
+        prompt_embeds_batches = torch.stack(prompt_embeds_batches)
+        pooled_prompt_embeds_batches = torch.stack(pooled_prompt_embeds_batches)
+        if text_ids is not None and text_ids.ndim == 3:
+            text_ids_batches = torch.stack(text_ids_batches)
+        if latent_image_ids is not None and latent_image_ids.ndim == 3:
+            latent_image_ids_batches = torch.stack(latent_image_ids_batches)
+        guidance_batches = torch.stack(guidance_batches)
+
+        return (
+            latents_batches,
+            prompt_embeds_batches,
+            pooled_prompt_embeds_batches,
+            text_ids_batches if text_ids.ndim == 3 else text_ids,
+            latent_image_ids_batches if latent_image_ids.ndim == 3 else latent_image_ids,
+            guidance_batches,
+            num_dummy_samples,
+        )
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        timesteps: List[int] = None,
+        guidance_scale: float = 3.5,
+        batch_size: int = 1,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        profiling_warmup_steps: Optional[int] = 0,
+        profiling_steps: Optional[int] = 0,
+        **kwargs,
+    ):
+        r"""
+        Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux.py#L531
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+            profiling_warmup_steps (`int`, *optional*):
+                Number of steps to ignore for profling.
+            profiling_steps (`int`, *optional*):
+                Number of steps to be captured when enabling profiling.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+        import habana_frameworks.torch as ht
+        import habana_frameworks.torch.core as htcore
+
+        quant_mode = kwargs["quant_mode"]
+
+        if quant_mode == "quantize-mixed":
+            import copy
+
+            transformer_bf16 = copy.deepcopy(self.transformer).to(self._execution_device)
+
+        if quant_mode == "measure" or quant_mode.startswith("quantize"):
+            import os
+
+            quant_config_path = os.getenv("QUANT_CONFIG")
+
+            htcore.hpu_set_env()
+
+            from neural_compressor.torch.quantization import FP8Config, convert, prepare
+
+            config = FP8Config.from_json_file(quant_config_path)
+            if config.measure:
+                self.transformer = prepare(self.transformer, config)
+            elif config.quantize:
+                self.transformer = convert(self.transformer, config)
+            htcore.hpu_initialize(self.transformer, mark_only_scales_as_const=True)
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            num_prompts = 1
+        elif prompt is not None and isinstance(prompt, list):
+            num_prompts = len(prompt)
+        else:
+            num_prompts = prompt_embeds.shape[0]
+        num_batches = math.ceil((num_images_per_prompt * num_prompts) / batch_size)
+
+        device = self._execution_device
+
+        # 3. Run text encoder
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            num_prompts * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        logger.info(
+            f"{num_prompts} prompt(s) received, {num_images_per_prompt} generation(s) per prompt,"
+            f" {batch_size} sample(s) per batch, {num_batches} total batch(es)."
+        )
+        if num_batches < 3:
+            logger.warning("The first two iterations are slower so it is recommended to feed more batches.")
+
+        throughput_warmup_steps = kwargs.get("throughput_warmup_steps", 3)
+        use_warmup_inference_steps = (
+            num_batches <= throughput_warmup_steps and num_inference_steps > throughput_warmup_steps
+        )
+
+        ht.hpu.synchronize()
+        t0 = time.time()
+        t1 = t0
+
+        hb_profiler = HabanaProfile(
+            warmup=profiling_warmup_steps,
+            active=profiling_steps,
+            record_shapes=False,
+        )
+        hb_profiler.start()
+
+        # 5.1. Split Input data to batches (HPU-specific step)
+        (
+            latents_batches,
+            text_embeddings_batches,
+            pooled_prompt_embeddings_batches,
+            text_ids_batches,
+            latent_image_ids_batches,
+            guidance_batches,
+            num_dummy_samples,
+        ) = self._split_inputs_into_batches(
+            batch_size, latents, prompt_embeds, pooled_prompt_embeds, text_ids, latent_image_ids, guidance
+        )
+
+        outputs = {
+            "images": [],
+        }
+
+        # 6. Denoising loop
+        for j in range(num_batches):
+            # The throughput is calculated from the 4th iteration
+            # because compilation occurs in the first 2-3 iterations
+            if j == throughput_warmup_steps:
+                ht.hpu.synchronize()
+                t1 = time.time()
+
+            latents_batch = latents_batches[0]
+            latents_batches = torch.roll(latents_batches, shifts=-1, dims=0)
+            text_embeddings_batch = text_embeddings_batches[0]
+            text_embeddings_batches = torch.roll(text_embeddings_batches, shifts=-1, dims=0)
+            pooled_prompt_embeddings_batch = pooled_prompt_embeddings_batches[0]
+            pooled_prompt_embeddings_batches = torch.roll(pooled_prompt_embeddings_batches, shifts=-1, dims=0)
+            if text_ids.ndim == 3:
+                text_ids_batch = text_ids_batches[0]
+                text_ids_batches = torch.roll(text_ids_batches, shifts=-1, dims=0)
+            if latent_image_ids.ndim == 3:
+                latent_image_ids_batch = latent_image_ids_batches[0]
+                latent_image_ids_batches = torch.roll(latent_image_ids_batches, shifts=-1, dims=0)
+            guidance_batch = None if guidance_batches[0].isnan().any() else guidance_batches[0]
+            guidance_batches = torch.roll(guidance_batches, shifts=-1, dims=0)
+
+            if hasattr(self.scheduler, "_init_step_index"):
+                # Reset scheduler step index for next batch
+                self.scheduler.timesteps = timesteps
+                self.scheduler._init_step_index(timesteps[0])
+
+            # Mixed quantization
+            quant_mixed_step = len(timesteps)
+            if quant_mode == "quantize-mixed":
+                # 10% of steps use higher precision in mixed quant mode
+                quant_mixed_step = quant_mixed_step - (quant_mixed_step // 10)
+                print(f"Use FP8  Transformer at steps 0 to {quant_mixed_step - 1}")
+                print(f"Use BF16 Transformer at steps {quant_mixed_step} to {len(timesteps) - 1}")
+
+            for i in self.progress_bar(range(len(timesteps))):
+                if use_warmup_inference_steps and i == throughput_warmup_steps and j == num_batches - 1:
+                    ht.hpu.synchronize()
+                    t1 = time.time()
+
+                if self.interrupt:
+                    continue
+
+                timestep = timesteps[0]
+                timesteps = torch.roll(timesteps, shifts=-1, dims=0)
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = timestep.expand(latents_batch.shape[0]).to(latents_batch.dtype)
+
+                if i >= quant_mixed_step:
+                    # Mixed quantization
+                    noise_pred = transformer_bf16(
+                        hidden_states=latents_batch,
+                        timestep=timestep / 1000,
+                        guidance=guidance_batch,
+                        pooled_projections=pooled_prompt_embeddings_batch,
+                        encoder_hidden_states=text_embeddings_batch,
+                        txt_ids=text_ids_batch if text_ids.ndim == 3 else text_ids,
+                        img_ids=latent_image_ids_batch if latent_image_ids.ndim ==3 else latent_image_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                else:
+                    noise_pred = self.transformer(
+                        hidden_states=latents_batch,
+                        timestep=timestep / 1000,
+                        guidance=guidance_batch,
+                        pooled_projections=pooled_prompt_embeddings_batch,
+                        encoder_hidden_states=text_embeddings_batch,
+                        txt_ids=text_ids_batch if text_ids.ndim == 3 else text_ids,
+                        img_ids=latent_image_ids_batch if latent_image_ids.ndim == 3 else latent_image_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                    )[0]
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents_batch.dtype
+                latents_batch = self.scheduler.step(noise_pred, timestep, latents_batch, return_dict=False)[0]
+
+                hb_profiler.step()
+                # htcore.mark_step(sync=True)
+
+            if not output_type == "latent":
+                latents_batch = self._unpack_latents(latents_batch, height, width, self.vae_scale_factor)
+                latents_batch = (latents_batch / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+                image = self.vae.decode(latents_batch, return_dict=False)[0]
+                image = self.image_processor.postprocess(image, output_type=output_type)
+            else:
+                image = latents_batch
+
+            outputs["images"].append(image)
+            # htcore.mark_step(sync=True)
+
+        # 7. Stage after denoising
+        hb_profiler.stop()
+
+        if quant_mode == "measure":
+            from neural_compressor.torch.quantization import finalize_calibration
+
+            finalize_calibration(self.transformer)
+
+        ht.hpu.synchronize()
+        speed_metrics_prefix = "generation"
+        if use_warmup_inference_steps:
+            t1 = warmup_inference_steps_time_adjustment(t1, t1, num_inference_steps, throughput_warmup_steps)
+        speed_measures = speed_metrics(
+            split=speed_metrics_prefix,
+            start_time=t0,
+            num_samples=batch_size
+            if t1 == t0 or use_warmup_inference_steps
+            else (num_batches - throughput_warmup_steps) * batch_size,
+            num_steps=batch_size * num_inference_steps
+            if use_warmup_inference_steps
+            else (num_batches - throughput_warmup_steps) * batch_size * num_inference_steps,
+            start_time_after_warmup=t1,
+        )
+        logger.info(f"Speed metrics: {speed_measures}")
+
+        # 8 Output Images
+        if num_dummy_samples > 0:
+            # Remove dummy generations if needed
+            outputs["images"][-1] = outputs["images"][-1][:-num_dummy_samples]
+
+        # Process generated images
+        for i, image in enumerate(outputs["images"][:]):
+            if i == 0:
+                outputs["images"].clear()
+
+            if output_type == "pil" and isinstance(image, list):
+                outputs["images"] += image
+            elif output_type in ["np", "numpy"] and isinstance(image, np.ndarray):
+                if len(outputs["images"]) == 0:
+                    outputs["images"] = image
+                else:
+                    outputs["images"] = np.concatenate((outputs["images"], image), axis=0)
+            else:
+                if len(outputs["images"]) == 0:
+                    outputs["images"] = image
+                else:
+                    outputs["images"] = torch.cat((outputs["images"], image), 0)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return outputs["images"]
+
+        return GaudiFluxPipelineOutput(
+            images=outputs["images"],
+            throughput=speed_measures[f"{speed_metrics_prefix}_samples_per_second"],
+        )
diff --git a/optimum/habana/diffusers/pipelines/pipeline_utils.py b/optimum/habana/diffusers/pipelines/pipeline_utils.py
index 7f36b90ae4..6e659edff4 100644
--- a/optimum/habana/diffusers/pipelines/pipeline_utils.py
+++ b/optimum/habana/diffusers/pipelines/pipeline_utils.py
@@ -55,6 +55,7 @@
     "optimum.habana.diffusers.schedulers": {
         "GaudiDDIMScheduler": ["save_pretrained", "from_pretrained"],
         "GaudiEulerDiscreteScheduler": ["save_pretrained", "from_pretrained"],
+        "GaudiFlowMatchEulerDiscreteScheduler": ["save_pretrained", "from_pretrained"],
         "GaudiEulerAncestralDiscreteScheduler": ["save_pretrained", "from_pretrained"],
     },
 }
diff --git a/optimum/habana/diffusers/schedulers/__init__.py b/optimum/habana/diffusers/schedulers/__init__.py
index 37eb80b1a6..48bf0bd8e9 100644
--- a/optimum/habana/diffusers/schedulers/__init__.py
+++ b/optimum/habana/diffusers/schedulers/__init__.py
@@ -1,3 +1,4 @@
 from .scheduling_ddim import GaudiDDIMScheduler
 from .scheduling_euler_ancestral_discrete import GaudiEulerAncestralDiscreteScheduler
 from .scheduling_euler_discrete import GaudiEulerDiscreteScheduler
+from .scheduling_flow_mactch_euler_discrete import GaudiFlowMatchEulerDiscreteScheduler
diff --git a/optimum/habana/diffusers/schedulers/scheduling_flow_mactch_euler_discrete.py b/optimum/habana/diffusers/schedulers/scheduling_flow_mactch_euler_discrete.py
new file mode 100644
index 0000000000..ccc597fa07
--- /dev/null
+++ b/optimum/habana/diffusers/schedulers/scheduling_flow_mactch_euler_discrete.py
@@ -0,0 +1,25 @@
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+
+
+class GaudiFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
+    # TODO: overwrite orginal func with following one to fix dyn error in gaudi lazy mode
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        # indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        # pos = 1 if len(indices) > 1 else 0
+
+        # return indices[pos].item()
+
+        masked = (schedule_timesteps == timestep)
+        tmp = masked.cumsum(dim=0)
+        pos = (tmp == 0).sum().item()
+        if masked.sum() > 1:
+            pos += (tmp == 1).sum().item()
+        return pos
diff --git a/optimum/habana/utils.py b/optimum/habana/utils.py
index 3da3f11872..5434344246 100755
--- a/optimum/habana/utils.py
+++ b/optimum/habana/utils.py
@@ -285,6 +285,7 @@ def __init__(
         warmup: int = 0,
         active: int = 0,
         record_shapes: bool = True,
+        with_stack: bool = False,
         output_dir: str = "./hpu_profile",
         wait: int = 0,
     ):
@@ -306,7 +307,7 @@ def noop():
                 activities=activities,
                 on_trace_ready=torch.profiler.tensorboard_trace_handler(output_dir),
                 record_shapes=record_shapes,
-                with_stack=False,
+                with_stack=with_stack
             )
             self.start = profiler.start
             self.stop = profiler.stop
diff --git a/setup.py b/setup.py
index ead947d01c..e8e693a108 100644
--- a/setup.py
+++ b/setup.py
@@ -33,7 +33,7 @@
     "optimum",
     "torch",
     "accelerate >= 0.33.0, < 0.34.0",
-    "diffusers == 0.29.2",
+    "diffusers == 0.31.0",
     "huggingface_hub >= 0.24.7",
     "sentence-transformers[train] == 3.0.1",
 ]