add experiments

app.py

@@ -29,16 +29,15 @@ from qwenimage.models.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-exp = ExperimentRegistry.get("qwen_lightning_lora")()
-
-exp.load()
 
 @spaces.GPU(duration=1500)
 def optim_pipe(exp):
+    exp = ExperimentRegistry.get("qwen_lightning_fa3_aot_int8_fuse_2step")()
+    exp.load()
     exp.optimize()
     return exp.pipe
 
-pipe = optim_pipe(exp)
+pipe = optim_pipe()
 
 
 MAX_SEED = np.iinfo(np.int32).max

qwenimage/experiments/experiments_qwen.py

@@ -255,14 +255,14 @@ class Qwen_Base_3step(QwenBaseExperiment):
         return self.pipe(*args, **kwargs).images[0]
 
 @ExperimentRegistry.register(name="qwen_lightning_lora_2step")
-class Qwen_Lightning_Lora_3step(Qwen_Lightning_Lora):
+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_3step(QwenBaseExperiment):
+class Qwen_Base_2step(QwenBaseExperiment):
     @ftimed
     def run_once(self, *args, **kwargs):
         kwargs["num_inference_steps"] = 2
@@ -338,7 +338,7 @@ class Qwen_FA3_AoT_int8(QwenBaseExperiment):
         )
 
 
-# @ExperimentRegistry.register(name="qwen_fp8")
+@ExperimentRegistry.register(name="qwen_fp8")
 class Qwen_fp8(QwenBaseExperiment):
     @ftimed
     def optimize(self):
@@ -347,7 +347,7 @@ class Qwen_fp8(QwenBaseExperiment):
         quantize_(self.pipe.transformer, Float8WeightOnlyConfig())
 
 
-# @ExperimentRegistry.register(name="qwen_int8")
+@ExperimentRegistry.register(name="qwen_int8")
 class Qwen_int8(QwenBaseExperiment):
     @ftimed
     def optimize(self):
@@ -410,3 +410,165 @@ class Qwen_FA3_AoT_fp8(QwenBaseExperiment):
 
         aoti_apply(compiled_transformer, self.pipe.transformer)
 
+# FA3_AoT_fp8_fuse
+@ExperimentRegistry.register(name="qwen_fa3_aot_fp8_fuse")
+class Qwen_FA3_AoT_fp8_fuse(QwenBaseExperiment):
+    @ftimed
+    # @spaces.GPU()
+    def optimize(self):
+        self.pipe.transformer.__class__ = QwenImageTransformer2DModel
+        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="_fa3_fuse"
+
+        cache_compiled=self.config.cache_compiled
+        
+        transformer_pt2_cache_path = f"checkpoints/transformer_fp8{suffix}_archive.pt2"
+        transformer_weights_cache_path = f"checkpoints/transformer_fp8{suffix}_weights.pt"
+
+        print(f"original model size: {get_model_size_in_bytes(self.pipe.transformer) / 1024 / 1024} MB")
+        quantize_(self.pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+        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)
+
+
+
+# FA3_AoT_int8_fuse
+@ExperimentRegistry.register(name="qwen_fa3_aot_int8_fuse")
+class Qwen_FA3_AoT_int8_fuse(QwenBaseExperiment):
+    @ftimed
+    def optimize(self):
+        self.pipe.transformer.__class__ = QwenImageTransformer2DModel
+        self.pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+        self.pipe.transformer.fuse_qkv_projections()
+        optimize_pipeline_(
+            self.pipe,
+            cache_compiled=self.config.cache_compiled,
+            quantize=True,
+            suffix="_fa3_fuse",
+            pipe_kwargs={
+                "image": [Image.new("RGB", (1024, 1024))],
+                "prompt":"prompt",
+                "num_inference_steps":4
+            }
+        )
+
+# lightning_FA3_AoT_fp8_fuse
+
+@ExperimentRegistry.register(name="qwen_lightning_fa3_aot_fp8_fuse")
+class Qwen_lightning_FA3_AoT_fp8_fuse(Qwen_Lightning_Lora):
+    @ftimed
+    # @spaces.GPU()
+    def optimize(self):
+        self.pipe.transformer.__class__ = QwenImageTransformer2DModel
+        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="_fa3_fuse"
+
+        cache_compiled=self.config.cache_compiled
+        
+        transformer_pt2_cache_path = f"checkpoints/transformer_fp8{suffix}_archive.pt2"
+        transformer_weights_cache_path = f"checkpoints/transformer_fp8{suffix}_weights.pt"
+
+        print(f"original model size: {get_model_size_in_bytes(self.pipe.transformer) / 1024 / 1024} MB")
+        quantize_(self.pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+        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)
+
+
+# lightning_FA3_AoT_int8_fuse
+
+@ExperimentRegistry.register(name="qwen_lightning_fa3_aot_int8_fuse")
+class Qwen_Lightning_FA3_AoT_int8_fuse(Qwen_Lightning_Lora):
+    @ftimed
+    def optimize(self):
+        self.pipe.transformer.__class__ = QwenImageTransformer2DModel
+        self.pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+        self.pipe.transformer.fuse_qkv_projections()
+        optimize_pipeline_(
+            self.pipe,
+            cache_compiled=self.config.cache_compiled,
+            quantize=True,
+            suffix="_fa3_fuse",
+            pipe_kwargs={
+                "image": [Image.new("RGB", (1024, 1024))],
+                "prompt":"prompt",
+                "num_inference_steps":4
+            }
+        )
+
+
+@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]

scripts/visual_compare.ipynb

@@ -2,25 +2,17 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "e5649df3",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "/home/ubuntu/Qwen-Image-Edit-Angles\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "%cd /home/ubuntu/Qwen-Image-Edit-Angles"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -30,53 +22,43 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/ubuntu/.local/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n",
-      "/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4\n",
-      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n",
-      "Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.14.1             Please see https://github.com/pytorch/ao/issues/2919 for more info\n",
-      "TMA benchmarks will be running without grid constant TMA descriptor.\n",
-      "2025-11-13 17:17:45.758469: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
-      "2025-11-13 17:17:45.772437: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
-      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
-      "E0000 00:00:1763054265.789504 1013472 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
-      "E0000 00:00:1763054265.795008 1013472 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
-      "W0000 00:00:1763054265.808263 1013472 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
-      "W0000 00:00:1763054265.808277 1013472 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
-      "W0000 00:00:1763054265.808280 1013472 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
-      "W0000 00:00:1763054265.808281 1013472 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
-      "2025-11-13 17:17:45.812477: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
-      "To enable the following instructions: AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
-      "Fetching 7 files: 100%|██████████| 7/7 [00:00<00:00, 81555.91it/s]\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from qwenimage.experiment import ExperimentConfig\n",
     "from qwenimage.experiments.experiments_qwen import ExperimentRegistry\n",
     "\n",
     "\n",
     "# experiment_names = ExperimentRegistry.keys()\n",
+    "# experiment_names = [\n",
+    "#     \"qwen_base\",\n",
+    "#     # \"qwen_fa3\",\n",
+    "#     # \"qwen_aot\",\n",
+    "#     # \"qwen_fa3_aot\",\n",
+    "#     # \"qwen_fa3_aot_int8\",\n",
+    "#     # \"qwen_fa3_aot_fp8\",\n",
+    "#     # \"qwen_fuse\",\n",
+    "#     \"qwen_base_3step\",\n",
+    "#     \"qwen_base_2step\",\n",
+    "#     \"qwen_lightning_lora\",\n",
+    "#     \"qwen_lightning_lora_3step\",\n",
+    "#     \"qwen_lightning_lora_2step\",\n",
+    "# ]\n",
+    "\n",
     "experiment_names = [\n",
-    "    \"qwen_base\",\n",
     "    # \"qwen_fa3\",\n",
     "    # \"qwen_aot\",\n",
     "    # \"qwen_fa3_aot\",\n",
-    "    # \"qwen_fa3_aot_int8\",\n",
-    "    # \"qwen_fa3_aot_fp8\",\n",
-    "    # \"qwen_fuse\",\n",
-    "    \"qwen_base_3step\",\n",
-    "    \"qwen_base_2step\",\n",
-    "    \"qwen_lightning_lora\",\n",
-    "    \"qwen_lightning_lora_3step\",\n",
-    "    \"qwen_lightning_lora_2step\",\n",
+    "    \"qwen_fa3_aot_int8\",\n",
+    "    \"qwen_fa3_aot_int8_fuse\",\n",
+    "    \"qwen_lightning_fa3_aot_int8_fuse\",\n",
+    "    # \"qwen_fa3_fuse\",\n",
+    "    \"qwen_fa3_aot_fp8\",\n",
+    "    \"qwen_fa3_aot_fp8_fuse\",\n",
+    "    \"qwen_lightning_fa3_aot_fp8_fuse\",\n",
+    "    \"qwen_lightning_fa3_aot_int8_fuse_2step\",\n",
+    "    \"qwen_base\",\n",
     "]\n",
     "\n",
     "report_dir = ExperimentConfig().report_dir\n",
@@ -94,21 +76,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": null,
    "id": "29077eb8",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "PosixPath('reports/qwen_base_outputs/000.jpg')"
-      ]
-     },
-     "execution_count": 23,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "experiment_outputs[\"qwen_base\"][0]"
    ]
@@ -119,119 +90,7 @@
    "id": "a591fdd6",
    "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": []
   },
   {
    "cell_type": "code",
@@ -249,17 +108,7 @@
    "execution_count": null,
    "id": "601aa246",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n",
-      "Processing rFID for qwen_base...\n",
-      "Calculating rFID: 4-step vs 3-step...\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "\n"
    ]
@@ -302,58 +151,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "41e36dc8",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Saved comparison grid 1/32 to reports/comparison_grid/comparison_000.jpg\n",
-      "Saved comparison grid 2/32 to reports/comparison_grid/comparison_001.jpg\n",
-      "Saved comparison grid 3/32 to reports/comparison_grid/comparison_002.jpg\n",
-      "Saved comparison grid 4/32 to reports/comparison_grid/comparison_003.jpg\n",
-      "Saved comparison grid 5/32 to reports/comparison_grid/comparison_004.jpg\n",
-      "Saved comparison grid 6/32 to reports/comparison_grid/comparison_005.jpg\n",
-      "Saved comparison grid 7/32 to reports/comparison_grid/comparison_006.jpg\n",
-      "Saved comparison grid 8/32 to reports/comparison_grid/comparison_007.jpg\n",
-      "Saved comparison grid 9/32 to reports/comparison_grid/comparison_008.jpg\n",
-      "Saved comparison grid 10/32 to reports/comparison_grid/comparison_009.jpg\n",
-      "Saved comparison grid 11/32 to reports/comparison_grid/comparison_010.jpg\n",
-      "Saved comparison grid 12/32 to reports/comparison_grid/comparison_011.jpg\n",
-      "Saved comparison grid 13/32 to reports/comparison_grid/comparison_012.jpg\n",
-      "Saved comparison grid 14/32 to reports/comparison_grid/comparison_013.jpg\n",
-      "Saved comparison grid 15/32 to reports/comparison_grid/comparison_014.jpg\n",
-      "Saved comparison grid 16/32 to reports/comparison_grid/comparison_015.jpg\n",
-      "Saved comparison grid 17/32 to reports/comparison_grid/comparison_016.jpg\n",
-      "Saved comparison grid 18/32 to reports/comparison_grid/comparison_017.jpg\n",
-      "Saved comparison grid 19/32 to reports/comparison_grid/comparison_018.jpg\n",
-      "Saved comparison grid 20/32 to reports/comparison_grid/comparison_019.jpg\n",
-      "Saved comparison grid 21/32 to reports/comparison_grid/comparison_020.jpg\n",
-      "Saved comparison grid 22/32 to reports/comparison_grid/comparison_021.jpg\n",
-      "Saved comparison grid 23/32 to reports/comparison_grid/comparison_022.jpg\n",
-      "Saved comparison grid 24/32 to reports/comparison_grid/comparison_023.jpg\n",
-      "Saved comparison grid 25/32 to reports/comparison_grid/comparison_024.jpg\n",
-      "Saved comparison grid 26/32 to reports/comparison_grid/comparison_025.jpg\n",
-      "Saved comparison grid 27/32 to reports/comparison_grid/comparison_026.jpg\n",
-      "Saved comparison grid 28/32 to reports/comparison_grid/comparison_027.jpg\n",
-      "Saved comparison grid 29/32 to reports/comparison_grid/comparison_028.jpg\n",
-      "Saved comparison grid 30/32 to reports/comparison_grid/comparison_029.jpg\n",
-      "Saved comparison grid 31/32 to reports/comparison_grid/comparison_030.jpg\n",
-      "Saved comparison grid 32/32 to reports/comparison_grid/comparison_031.jpg\n",
-      "\n",
-      "All comparison grids saved to reports/comparison_grid\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import math\n",
     "from PIL import Image\n",
     "import numpy as np\n",
     "from pathlib import Path\n",
     "\n",
-    "comparison_dir = report_dir / \"comparison_grid\"\n",
+    "comparison_dir = report_dir / \"comparison_grid_2\"\n",
     "comparison_dir.mkdir(exist_ok=True, parents=True)\n",
     "\n",
     "\n",
@@ -423,7 +231,119 @@
    "id": "244dfe0f",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "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"
+   ]
   }
  ],
  "metadata": {