feat: Refactor AOT optimization in pipeline to enhance dynamic shape handling and error management

optimization.py

@@ -1,57 +1,77 @@
-from typing import Any
-from typing import Callable
-from typing import ParamSpec
-from torchao.quantization import quantize_
-from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
+# optimization.py
+from typing import Any, Callable, ParamSpec
 import spaces
 import torch
 from torch.utils._pytree import tree_map
 
-P = ParamSpec('P')
+# No usamos torchao ni quantización, para evitar conflictos con versiones
+P = ParamSpec("P")
 
-TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim('image_seq_length')
-TRANSFORMER_TEXT_SEQ_LENGTH_DIM = torch.export.Dim('text_seq_length')
+# Definimos dimensiones dinámicas que sí existen en tu pipeline Qwen-Image
+TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim("image_seq_length")
+TRANSFORMER_TEXT_SEQ_LENGTH_DIM = torch.export.Dim("text_seq_length")
 
+# Solo incluimos las claves que realmente aparecen en tu modelo
 TRANSFORMER_DYNAMIC_SHAPES = {
-    'hidden_states': {
-        1: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
-    },
-    'encoder_hidden_states': {
-        1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
-    },
-    'encoder_hidden_states_mask': {
-        1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
-    },
-    'image_rotary_emb': ({
-        0: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
-    }, {
-        0: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
-    }),
+    "hidden_states": {1: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM},
+    "encoder_hidden_states": {1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM},
+    "encoder_hidden_states_mask": {1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM},
 }
 
+# Configuraciones del compilador AOTInductor
 INDUCTOR_CONFIGS = {
-    'conv_1x1_as_mm': True,
-    'epilogue_fusion': False,
-    'coordinate_descent_tuning': True,
-    'coordinate_descent_check_all_directions': True,
-    'max_autotune': True,
-    'triton.cudagraphs': True,
+    "conv_1x1_as_mm": True,
+    "epilogue_fusion": False,
+    "coordinate_descent_tuning": True,
+    "coordinate_descent_check_all_directions": True,
+    "max_autotune": True,
+    "triton.cudagraphs": True,
 }
 
+
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
-    @spaces.GPU(duration=1500)
-    def compile_transformer():
-        with spaces.aoti_capture(pipeline.transformer) as call:
-            pipeline(*args, **kwargs)
-        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
-        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
-        # quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
-        exported = torch.export.export(
-            mod=pipeline.transformer,
-            args=call.args,
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes,
-        )
-        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
-
-    spaces.aoti_apply(compile_transformer(), pipeline.transformer)
+    """
+    Optimiza el transformer interno del pipeline de Qwen-Image usando AOTInductor.
+    Funciona con ZeroGPU y solo compila el módulo transformer, no todo el pipeline.
+    """
+
+    # Si no hay GPU, no intentamos compilar
+    if not torch.cuda.is_available():
+        print("⚠️ CUDA no disponible. Se omite AOT.")
+        return pipeline
+
+    try:
+        @spaces.GPU(duration=1200)
+        def compile_transformer():
+            print("🏗️ Capturando modelo para AOT...")
+
+            # Capturamos la llamada del transformer
+            with spaces.aoti_capture(pipeline.transformer) as call:
+                pipeline(*args, **kwargs)
+
+            # Creamos el mapa de shapes dinámicos solo con claves válidas
+            dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+            # Añadimos los shapes esperados, filtrando los que realmente existen
+            for k, v in TRANSFORMER_DYNAMIC_SHAPES.items():
+                if k in call.kwargs:
+                    dynamic_shapes[k] = v
+
+            print("🚀 Exportando modelo con torch.export...")
+            exported = torch.export.export(
+                mod=pipeline.transformer,
+                args=call.args,
+                kwargs=call.kwargs,
+                dynamic_shapes=dynamic_shapes,
+            )
+
+            print("⚙️ Compilando con AOTInductor...")
+            return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
+
+        print("🧠 Aplicando AOT al transformer...")
+        spaces.aoti_apply(compile_transformer(), pipeline.transformer)
+        print("✅ AOT aplicado correctamente al transformer de Qwen-Image.")
+
+    except Exception as e:
+        print(f"⚠️ Error al aplicar AOT: {e}")
+
+    return pipeline