Update README.md

README.md

@@ -38,7 +38,7 @@ Example: `python -m deepcompressor.app.diffusion.ptq examples/diffusion/configs/
 
 Folder Structure 
 
-- refer [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)
+- refer [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev/tree/main)
 
 - refer [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev/tree/main)
 
@@ -51,11 +51,7 @@ potential fix: app.diffusion.pipeline.config.py
 ```python
     @staticmethod
     def _default_build(
-        name: str,
-        path: str,
-        dtype: str | torch.dtype,
-        device: str | torch.device,
-        shift_activations: bool
+        name: str, path: str, dtype: str | torch.dtype, device: str | torch.device, shift_activations: bool
     ) -> DiffusionPipeline:
         if not path:
             if name == "sdxl":
@@ -64,6 +60,8 @@ potential fix: app.diffusion.pipeline.config.py
                 path = "stabilityai/sdxl-turbo"
             elif name == "pixart-sigma":
                 path = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
+            elif name == "flux.1-kontext-dev":
+                path = "black-forest-labs/FLUX.1-Kontext-dev"
             elif name == "flux.1-dev":
                 path = "black-forest-labs/FLUX.1-dev"
             elif name == "flux.1-canny-dev":
@@ -76,24 +74,22 @@ potential fix: app.diffusion.pipeline.config.py
                 path = "black-forest-labs/FLUX.1-schnell"
             else:
                 raise ValueError(f"Path for {name} is not specified.")
-    
-        # Instantiate the pipeline
-        if name in ["flux.1-canny-dev", "flux.1-depth-dev"]:
+        if name in ["flux.1-kontext-dev"]:
+            pipeline = FluxKontextPipeline.from_pretrained(path, torch_dtype=dtype)
+        elif name in ["flux.1-canny-dev", "flux.1-depth-dev"]:
             pipeline = FluxControlPipeline.from_pretrained(path, torch_dtype=dtype)
         elif name == "flux.1-fill-dev":
             pipeline = FluxFillPipeline.from_pretrained(path, torch_dtype=dtype)
         elif name.startswith("sana-"):
             if dtype == torch.bfloat16:
-                pipeline = SanaPipeline.from_pretrained(
-                    path, variant="bf16", torch_dtype=dtype, use_safetensors=True
-                )
+                pipeline = SanaPipeline.from_pretrained(path, variant="bf16", torch_dtype=dtype, use_safetensors=True)
                 pipeline.vae.to(dtype)
                 pipeline.text_encoder.to(dtype)
             else:
                 pipeline = SanaPipeline.from_pretrained(path, torch_dtype=dtype)
         else:
             pipeline = AutoPipelineForText2Image.from_pretrained(path, torch_dtype=dtype)
-    
+
         # Debug output
         print(">>> DEVICE:", device)
         print(">>> PIPELINE TYPE:", type(pipeline))
@@ -128,14 +124,14 @@ potential fix: app.diffusion.pipeline.config.py
 
 Debug Log
 ```
-25-07-21 22:47:02 | I | === Start Evaluating ===
-25-07-21 22:47:02 | I | * Building diffusion model pipeline
-Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 15.44it/s]
-Loading pipeline components...:  57%|████████████████████████████████████████████████████████████████                                            | 4/7 [00:00<00:00,  9.47it/s]
+25-07-22 20:11:56 | I | === Start Evaluating ===
+25-07-22 20:11:56 | I | * Building diffusion model pipeline
+Loading pipeline components...:   0%|                                                             | 0/7 [00:00<?, ?it/s]
 You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
-Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  7.79it/s]
+Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 18.92it/s]
+Loading pipeline components...: 100%|█████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.50it/s]
 >>> DEVICE: cuda
->>> PIPELINE TYPE: <class 'diffusers.pipelines.flux.pipeline_flux.FluxPipeline'>
+>>> PIPELINE TYPE: <class 'diffusers.pipelines.flux.pipeline_flux_kontext.FluxKontextPipeline'>
 >>> Moving transformer to cuda using to_empty()
 >>> WARNING: transformer.to_empty(cuda) failed: Module.to_empty() takes 1 positional argument but 2 were given
 >>> Falling back to transformer.to(cuda)
@@ -146,16 +142,16 @@ Loading pipeline components...: 100%|██████████████
 >>> Moving text_encoder to cuda using to_empty()
 >>> WARNING: text_encoder.to_empty(cuda) failed: Module.to_empty() takes 1 positional argument but 2 were given
 >>> Falling back to text_encoder.to(cuda)
-25-07-21 22:47:05 | I |   Replacing fused Linear with ConcatLinear.
-25-07-21 22:47:05 | I |     + Replacing fused Linear in single_transformer_blocks.0 with ConcatLinear.
-25-07-21 22:47:05 | I |       - in_features = 3072/15360
-25-07-21 22:47:05 | I |       - out_features = 3072
-25-07-21 22:47:05 | I |     + Replacing fused Linear in single_transformer_blocks.1 with ConcatLinear.
-25-07-21 22:47:05 | I |       - in_features = 3072/15360
-25-07-21 22:47:05 | I |       - out_features = 3072
-25-07-21 22:47:05 | I |     + Replacing fused Linear in single_transformer_blocks.2 with ConcatLinear.
-25-07-21 22:47:05 | I |       - in_features = 3072/15360
-25-07-21 22:47:05 | I |       - out_features = 3072
+25-07-22 20:11:59 | I |   Replacing fused Linear with ConcatLinear.
+25-07-22 20:11:59 | I |     + Replacing fused Linear in single_transformer_blocks.0 with ConcatLinear.
+25-07-22 20:11:59 | I |       - in_features = 3072/15360
+25-07-22 20:11:59 | I |       - out_features = 3072
+25-07-22 20:11:59 | I |     + Replacing fused Linear in single_transformer_blocks.1 with ConcatLinear.
+25-07-22 20:11:59 | I |       - in_features = 3072/15360
+25-07-22 20:11:59 | I |       - out_features = 3072
+25-07-22 20:11:59 | I |     + Replacing fused Linear in single_transformer_blocks.2 with ConcatLinear.
+25-07-22 20:11:59 | I |       - in_features = 3072/15360
+25-07-22 20:11:59 | I |       - out_features = 3072
 ```
 
 2) KeyError: <class 'diffusers.models.transformers.transformer_flux.FluxAttention'>