add fbcache, lpips comparison, 50 step

qwenimage/experiments/experiments_qwen.py

@@ -12,6 +12,7 @@ import torch
 from PIL import Image
 import pandas as pd
 from spaces.zero.torch.aoti import ZeroGPUCompiledModel, ZeroGPUWeights
+from torchao import autoquant
 from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, Float8WeightOnlyConfig, Int4WeightOnlyConfig, Int8DynamicActivationInt4WeightConfig, Int8DynamicActivationInt8WeightConfig, quantize_
 from torchao.quantization import Int8WeightOnlyConfig
 import spaces
@@ -19,7 +20,8 @@ import torch
 from torch.utils._pytree import tree_map
 from torchao.utils import get_model_size_in_bytes
 
-from qwenimage.debug import ftimed, print_first_param
+from qwenimage.debug import ctimed, ftimed, print_first_param
+from qwenimage.models.first_block_cache import apply_cache_on_pipe
 from qwenimage.models.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.models.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.models.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
@@ -48,6 +50,19 @@ class ExperimentRegistry:
             raise KeyError(f"{name} not in {list(cls.registry.keys())}")
         return cls.registry[name]
     
+    @classmethod
+    def filter(cls, startswith=None, endswith=None, contains=None, not_contain=None):
+        keys = list(cls.registry.keys())
+        if startswith is not None:
+            keys = [k for k in keys if k.startswith(startswith)]
+        if endswith is not None:
+            keys = [k for k in keys if k.endswith(endswith)]
+        if contains is not None:
+            keys = [k for k in keys if contains in k]
+        if not_contain is not None:
+            keys = [k for k in keys if not_contain not in k]
+        return keys
+
     @classmethod
     def keys(cls):
         return list(cls.registry.keys())
@@ -142,7 +157,6 @@ class QwenBaseExperiment(AbstractExperiment):
     def optimize(self):
         pass
     
-    @ftimed
     def run_once(self, *args, **kwargs):
         return self.pipe(*args, **kwargs).images[0]
 
@@ -152,7 +166,8 @@ class QwenBaseExperiment(AbstractExperiment):
 
         for i in range(self.config.iterations):
             inputs = self.pipe_inputs[i]
-            output = self.run_once(**inputs)
+            with ctimed("run_once"):
+                output = self.run_once(**inputs)
             output.save(output_save_dir / f"{i:03d}.jpg")
 
     def report(self):
@@ -182,6 +197,40 @@ class QwenBaseExperiment(AbstractExperiment):
         del self.pipe.transformer
         del self.pipe
 
+@ExperimentRegistry.register(name="qwen_50step")
+class Qwen_50Step(QwenBaseExperiment):
+    @ftimed
+    def load(self):
+        dtype = torch.bfloat16
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"experiment load cuda: {torch.cuda.is_available()=}")
+
+        pipe = QwenImageEditPlusPipeline.from_pretrained(
+            "Qwen/Qwen-Image-Edit-2509", 
+            transformer=QwenImageTransformer2DModel.from_pretrained( # use our own model
+                "Qwen/Qwen-Image-Edit-2509",
+                subfolder='transformer',
+                torch_dtype=dtype,
+                device_map=device
+            ),
+            torch_dtype=dtype,
+        ).to(device)
+
+        pipe.load_lora_weights(
+            "dx8152/Qwen-Edit-2509-Multiple-angles", 
+            weight_name="镜头转换.safetensors", adapter_name="angles"
+        )
+
+        pipe.set_adapters(["angles"], adapter_weights=[1.])
+        pipe.fuse_lora(adapter_names=["angles"], lora_scale=1.25)
+        pipe.unload_lora_weights()
+        self.pipe = pipe
+
+    def run_once(self, *args, **kwargs):
+        kwargs["num_inference_steps"] = 50
+        return self.pipe(*args, **kwargs).images[0]
+
+
 @ExperimentRegistry.register(name="qwen_lightning_lora")
 class Qwen_Lightning_Lora(QwenBaseExperiment):
     @ftimed
@@ -243,28 +292,24 @@ class Qwen_Lightning_Lora(QwenBaseExperiment):
 
 @ExperimentRegistry.register(name="qwen_lightning_lora_3step")
 class Qwen_Lightning_Lora_3step(Qwen_Lightning_Lora):
-    @ftimed
     def run_once(self, *args, **kwargs):
         kwargs["num_inference_steps"] = 3
         return self.pipe(*args, **kwargs).images[0]
 
 @ExperimentRegistry.register(name="qwen_base_3step")
 class Qwen_Base_3step(QwenBaseExperiment):
-    @ftimed
     def run_once(self, *args, **kwargs):
         kwargs["num_inference_steps"] = 3
         return self.pipe(*args, **kwargs).images[0]
 
 @ExperimentRegistry.register(name="qwen_lightning_lora_2step")
 class Qwen_Lightning_Lora_2step(Qwen_Lightning_Lora):
-    @ftimed
     def run_once(self, *args, **kwargs):
         kwargs["num_inference_steps"] = 2
         return self.pipe(*args, **kwargs).images[0]
 
 @ExperimentRegistry.register(name="qwen_base_2step")
 class Qwen_Base_2step(QwenBaseExperiment):
-    @ftimed
     def run_once(self, *args, **kwargs):
         kwargs["num_inference_steps"] = 2
         return self.pipe(*args, **kwargs).images[0]
@@ -582,20 +627,129 @@ class Qwen_Lightning_FA3_AoT_int8_fuse(Qwen_Lightning_Lora):
 
 @ExperimentRegistry.register(name="qwen_lightning_fa3_aot_int8_fuse_2step")
 class Qwen_Lightning_FA3_AoT_int8_fuse_2step(Qwen_Lightning_FA3_AoT_int8_fuse):
-    @ftimed
     def run_once(self, *args, **kwargs):
         kwargs["num_inference_steps"] = 2
         return self.pipe(*args, **kwargs).images[0]
 
+@ExperimentRegistry.register(name="qwen_lightning_fa3_aot_int8_fuse_3step")
+class Qwen_Lightning_FA3_AoT_int8_fuse_3step(Qwen_Lightning_FA3_AoT_int8_fuse):
+    def run_once(self, *args, **kwargs):
+        kwargs["num_inference_steps"] = 3
+        return self.pipe(*args, **kwargs).images[0]
+
+@ExperimentRegistry.register(name="qwen_fa3_aot_int8_fuse_2step")
+class Qwen_FA3_AoT_int8_fuse_2step(Qwen_FA3_AoT_int8_fuse):
+    def run_once(self, *args, **kwargs):
+        kwargs["num_inference_steps"] = 2
+        return self.pipe(*args, **kwargs).images[0]
+
+@ExperimentRegistry.register(name="qwen_fa3_aot_int8_fuse_3step")
+class Qwen_FA3_AoT_int8_fuse_3step(Qwen_FA3_AoT_int8_fuse):
+    def run_once(self, *args, **kwargs):
+        kwargs["num_inference_steps"] = 3
+        return self.pipe(*args, **kwargs).images[0]
 
 @ExperimentRegistry.register(name="qwen_channels_last")
 class Qwen_Channels_Last(QwenBaseExperiment):
     """
-    This experiment is fully useless: channels last format only works with NCHW tensors, 
+    This experiment may be useless: channels last format only works with NCHW tensors, 
     i.e. 2D CNNs, transformer is 1D and vae is 3D, plus, for it to work the inputs need to 
     be converted in-pipe as well. left for reference.
     """
     @ftimed
     def optimize(self):
-        self.pipe.vae = self.pipe.vae.to(memory_format=torch.channels_last)
-        self.pipe.transformer = self.pipe.transformer.to(memory_format=torch.channels_last)
+        # self.pipe.vae = self.pipe.vae.to(memory_format=torch.channels_last_3d)
+        self.pipe.transformer = self.pipe.transformer.to(memory_format=torch.channels_last)
+
+@ExperimentRegistry.register(name="qwen_fbcache_05")
+class Qwen_FBCache_05(QwenBaseExperiment):
+    @ftimed
+    def optimize(self):
+        apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.5,)
+
+
+@ExperimentRegistry.register(name="qwen_fbcache_055")
+class Qwen_FBCache_055(QwenBaseExperiment):
+    @ftimed
+    def optimize(self):
+        apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.55,)
+
+@ExperimentRegistry.register(name="qwen_fbcache_054")
+class Qwen_FBCache_054(QwenBaseExperiment):
+    @ftimed
+    def optimize(self):
+        apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.54,)
+
+@ExperimentRegistry.register(name="qwen_fbcache_053")
+class Qwen_FBCache_053(QwenBaseExperiment):
+    @ftimed
+    def optimize(self):
+        apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.53,)
+
+@ExperimentRegistry.register(name="qwen_fbcache_052")
+class Qwen_FBCache_052(QwenBaseExperiment):
+    @ftimed
+    def optimize(self):
+        apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.52,)
+
+@ExperimentRegistry.register(name="qwen_fbcache_051")
+class Qwen_FBCache_051(QwenBaseExperiment):
+    @ftimed
+    def optimize(self):
+        apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.51,)
+
+
+# @ExperimentRegistry.register(name="qwen_lightning_fa3_aot_autoquant_fuse")
+class Qwen_lightning_FA3_AoT_autoquant_fuse(Qwen_Lightning_Lora):
+    """
+    Seemingly not working with AoT export
+    """
+    @ftimed
+    def optimize(self):
+        self.pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+        self.pipe.transformer.fuse_qkv_projections()
+
+        pipe_kwargs={
+            "image": [Image.new("RGB", (1024, 1024))],
+            "prompt":"prompt",
+            "num_inference_steps":4
+        }
+        suffix="_autoquant_fa3_fuse"
+
+        cache_compiled=self.config.cache_compiled
+        
+        transformer_pt2_cache_path = f"checkpoints/transformer_{suffix}_archive.pt2"
+        transformer_weights_cache_path = f"checkpoints/transformer_{suffix}_weights.pt"
+
+        print(f"original model size: {get_model_size_in_bytes(self.pipe.transformer) / 1024 / 1024} MB")
+        autoquant(self.pipe.transformer, error_on_unseen=False)
+        print_first_param(self.pipe.transformer)
+        print(f"quantized model size: {get_model_size_in_bytes(self.pipe.transformer) / 1024 / 1024} MB")
+
+        inductor_config = INDUCTOR_CONFIGS
+
+        if os.path.isfile(transformer_pt2_cache_path) and cache_compiled:
+            drain_module_parameters(self.pipe.transformer)
+            zerogpu_weights = torch.load(transformer_weights_cache_path, weights_only=False)
+            compiled_transformer = ZeroGPUCompiledModel(transformer_pt2_cache_path, zerogpu_weights)
+        else:
+            with spaces.aoti_capture(self.pipe.transformer) as call:
+                self.pipe(**pipe_kwargs)
+
+            dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+            dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
+            
+            exported = torch.export.export(
+                mod=self.pipe.transformer,
+                args=call.args,
+                kwargs=call.kwargs,
+                dynamic_shapes=dynamic_shapes,
+            )
+
+            compiled_transformer = spaces.aoti_compile(exported, inductor_config)
+            with open(transformer_pt2_cache_path, "wb") as f:
+                f.write(compiled_transformer.archive_file.getvalue())
+            torch.save(compiled_transformer.weights, transformer_weights_cache_path)
+
+
+        aoti_apply(compiled_transformer, self.pipe.transformer)

qwenimage/models/first_block_cache.py

@@ -0,0 +1,73 @@
+import functools
+import unittest
+
+import torch
+
+from qwenimage.models.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+from qwenimage.models.transformer_qwenimage import QwenImageTransformer2DModel
+
+from para_attn.first_block_cache import utils
+
+
+def apply_cache_on_transformer(
+    transformer: QwenImageTransformer2DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            utils.CachedTransformerBlocks(
+                transformer.transformer_blocks,
+                transformer=transformer,
+                return_hidden_states_first=False,
+            )
+        ]
+    )
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with unittest.mock.patch.object(
+            self,
+            "transformer_blocks",
+            cached_transformer_blocks,
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_cached = True
+
+    return transformer
+
+
+def apply_cache_on_pipe(
+    pipe: QwenImageEditPlusPipeline,
+    *,
+    shallow_patch: bool = False,
+    **kwargs,
+):
+    if not getattr(pipe, "_is_cached", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs_):
+            with utils.cache_context(utils.create_cache_context(**kwargs)):
+                return original_call(self, *args, **kwargs_)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_cached = True
+
+    if not shallow_patch:
+        apply_cache_on_transformer(pipe.transformer)
+
+    return pipe

qwenimage/models/pipeline_qwenimage_edit_plus.py

@@ -548,6 +548,7 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
         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,
+        channels_last_format: bool = False,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -665,6 +666,7 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
             self._attention_kwargs = attention_kwargs
             self._current_timestep = None
             self._interrupt = False
+            self.channels_last_format = channels_last_format
 
             # 2. Define call parameters
             if prompt is not None and isinstance(prompt, str):

qwenimage/models/transformer_qwenimage.py

@@ -15,6 +15,7 @@
 import functools
 import math
 from typing import Any, Dict, List, Optional, Tuple, Union
+import warnings
 
 import torch
 import torch.nn as nn
@@ -615,6 +616,7 @@ class QwenImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Fro
 
         for index_block, block in enumerate(self.transformer_blocks):
             if torch.is_grad_enabled() and self.gradient_checkpointing:
+                warnings.warn("Gradient ckpt?")
                 encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
                     block,
                     hidden_states,

requirements.txt

@@ -15,4 +15,5 @@ pydantic
 pandas
 modal
 
+para-attn
 lpips

scripts/scratch.ipynb

@@ -606,26 +606,7 @@
    "id": "e571d339",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "mvae_params = sum(p.numel() for p in flux.blend_model.mvae.parameters())\n",
-    "model_size_mb = sum(p.numel() * p.element_size() for p in flux.blend_model.mvae.parameters()) / (1024 ** 2)\n",
-    "print(f\"mvae parameter count: {mvae_params:,}\")\n",
-    "print(f\"mvae model size: {model_size_mb:.2f} MB\")\n",
-    "\n",
-    "transformer_params = sum(p.numel() for p in flux.transformer.parameters())\n",
-    "model_size_mb = sum(p.numel() * p.element_size() for p in flux.transformer.parameters()) / (1024 ** 2)\n",
-    "print(f\"flux.transformer parameter count: {transformer_params:,}\")\n",
-    "print(f\"flux.transformer model size: {model_size_mb:.2f} MB\")\n",
-    "\n",
-    "vae_params = sum(p.numel() for p in flux.vae.parameters())\n",
-    "model_size_mb = sum(p.numel() * p.element_size() for p in flux.vae.parameters()) / (1024 ** 2)\n",
-    "print(f\"flux.vae parameter count: {vae_params:,}\")\n",
-    "print(f\"flux.vae model size: {model_size_mb:.2f} MB\")\n",
-    "\n",
-    "print(f\"\\nParameter comparisons:\")\n",
-    "print(f\"mvae vs transformer: {mvae_params / transformer_params * 100:.3f}%\")\n",
-    "print(f\"mvae vs vae: {mvae_params / vae_params * 100:.3f}%\")\n"
-   ]
+   "source": []
   },
   {
    "cell_type": "code",

scripts/visual_compare.ipynb

@@ -231,119 +231,7 @@
    "id": "244dfe0f",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "\n",
-    "import lpips\n",
-    "import torch\n",
-    "from PIL import Image\n",
-    "import torchvision.transforms as transforms\n",
-    "\n",
-    "# Initialize LPIPS model\n",
-    "loss_fn = lpips.LPIPS(net='alex')  # or 'vgg' or 'squeeze'\n",
-    "if torch.cuda.is_available():\n",
-    "    loss_fn = loss_fn.cuda()\n",
-    "\n",
-    "# Transform to convert PIL images to tensors\n",
-    "transform = transforms.Compose([\n",
-    "    transforms.ToTensor(),\n",
-    "    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n",
-    "])\n",
-    "\n",
-    "def calculate_lpips_scores(base_paths, compare_paths):\n",
-    "    \"\"\"Calculate LPIPS scores between two sets of images.\"\"\"\n",
-    "    scores = []\n",
-    "    \n",
-    "    # Get the minimum number of images available\n",
-    "    num_images = min(len(base_paths), len(compare_paths))\n",
-    "    \n",
-    "    for idx in range(num_images):\n",
-    "        # Load images\n",
-    "        img1 = Image.open(base_paths[idx]).convert('RGB')\n",
-    "        img2 = Image.open(compare_paths[idx]).convert('RGB')\n",
-    "        \n",
-    "        # Resize if dimensions don't match\n",
-    "        if img1.size != img2.size:\n",
-    "            img2 = img2.resize(img1.size, Image.LANCZOS)\n",
-    "        \n",
-    "        # Transform to tensors\n",
-    "        img1_tensor = transform(img1).unsqueeze(0)\n",
-    "        img2_tensor = transform(img2).unsqueeze(0)\n",
-    "        \n",
-    "        if torch.cuda.is_available():\n",
-    "            img1_tensor = img1_tensor.cuda()\n",
-    "            img2_tensor = img2_tensor.cuda()\n",
-    "        \n",
-    "        # Calculate LPIPS\n",
-    "        with torch.no_grad():\n",
-    "            score = loss_fn(img1_tensor, img2_tensor)\n",
-    "        \n",
-    "        scores.append(score.item())\n",
-    "    \n",
-    "    return scores\n",
-    "\n",
-    "# Define experiment sets\n",
-    "experiment_sets = {\n",
-    "    'qwen_base': {\n",
-    "        '4step': 'qwen_base',\n",
-    "        '3step': 'qwen_base_3step',\n",
-    "        '2step': 'qwen_base_2step'\n",
-    "    },\n",
-    "    'qwen_lightning_lora': {\n",
-    "        '4step': 'qwen_lightning_lora',\n",
-    "        '3step': 'qwen_lightning_lora_3step',\n",
-    "        '2step': 'qwen_lightning_lora_2step'\n",
-    "    }\n",
-    "}\n",
-    "\n",
-    "# Calculate LPIPS scores for each set\n",
-    "results = {}\n",
-    "\n",
-    "for set_name, experiments in experiment_sets.items():\n",
-    "    print(f\"\\nProcessing {set_name}...\")\n",
-    "    \n",
-    "    # Get image paths\n",
-    "    base_4step_paths = experiment_outputs[experiments['4step']]\n",
-    "    step_3_paths = experiment_outputs[experiments['3step']]\n",
-    "    step_2_paths = experiment_outputs[experiments['2step']]\n",
-    "    \n",
-    "    # Calculate LPIPS scores\n",
-    "    print(f\"Calculating LPIPS: 4-step vs 3-step...\")\n",
-    "    scores_4vs3 = calculate_lpips_scores(base_4step_paths, step_3_paths)\n",
-    "    \n",
-    "    print(f\"Calculating LPIPS: 4-step vs 2-step...\")\n",
-    "    scores_4vs2 = calculate_lpips_scores(base_4step_paths, step_2_paths)\n",
-    "    \n",
-    "    # Create results dataframe\n",
-    "    results_df = pd.DataFrame({\n",
-    "        'comparison': ['4step_vs_3step', '4step_vs_2step'],\n",
-    "        'mean_lpips': [\n",
-    "            np.mean(scores_4vs3),\n",
-    "            np.mean(scores_4vs2)\n",
-    "        ],\n",
-    "        'std_lpips': [\n",
-    "            np.std(scores_4vs3),\n",
-    "            np.std(scores_4vs2)\n",
-    "        ],\n",
-    "        'num_samples': [\n",
-    "            len(scores_4vs3),\n",
-    "            len(scores_4vs2)\n",
-    "        ]\n",
-    "    })\n",
-    "    \n",
-    "    # Save to CSV\n",
-    "    csv_path = report_dir / f\"lpips_scores_{set_name}.csv\"\n",
-    "    results_df.to_csv(csv_path, index=False)\n",
-    "    \n",
-    "    print(f\"\\nResults for {set_name}:\")\n",
-    "    print(results_df)\n",
-    "    print(f\"\\nSaved to: {csv_path}\")\n",
-    "    \n",
-    "    results[set_name] = results_df\n",
-    "\n",
-    "print(\"\\n\" + \"=\"*60)\n",
-    "print(\"LPIPS Analysis Complete!\")\n",
-    "print(\"=\"*60)\n"
-   ]
+   "source": []
   }
  ],
  "metadata": {