feat: Refactor optimize_pipeline_ function for improved dynamic shape handling and stability in ZeroGPU

optimization.py

@@ -6,11 +6,9 @@ from torch.utils._pytree import tree_map
 
 P = ParamSpec("P")
 
-# Dimensiones fijas sugeridas por torch.export (evitan el UserError)
 TEXT_SEQ_LENGTH = 12
 IMAGE_SEQ_LENGTH = 4096
 
-# Configuraciones del compilador AOTInductor
 INDUCTOR_CONFIGS = {
     "conv_1x1_as_mm": True,
     "epilogue_fusion": False,
@@ -23,8 +21,10 @@ INDUCTOR_CONFIGS = {
 
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
     """
-    Compila el transformer de Qwen-Image usando AOTInductor para ZeroGPU.
-    Versión adaptada específicamente a tu pipeline: sin dynamic_shapes conflictivos.
+    Versión estable para Qwen-Image:
+    - Corrige estructura de dynamic_shapes (incluye img_shapes como lista)
+    - Usa longitudes fijas sugeridas
+    - Funciona en ZeroGPU sin romper el Space
     """
 
     if not torch.cuda.is_available():
@@ -36,23 +36,24 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
         def compile_transformer():
             print("🏗️ Capturando modelo para AOT...")
 
-            # Captura del grafo del transformer
             with spaces.aoti_capture(pipeline.transformer) as call:
                 pipeline(*args, **kwargs)
 
-            # Asignamos longitudes fijas a las secuencias (para evitar UserError)
+            # Base dinámica (misma estructura que inputs)
             dynamic_shapes = tree_map(lambda t: None, call.kwargs)
-            dynamic_shapes = {
-                k: None for k in call.kwargs.keys()
-            }  # evitamos conflicto de claves
+
+            # Ajustes estáticos seguros
             static_shapes = {
                 "hidden_states": {1: IMAGE_SEQ_LENGTH},
                 "encoder_hidden_states": {1: TEXT_SEQ_LENGTH},
                 "encoder_hidden_states_mask": {1: TEXT_SEQ_LENGTH},
+                # img_shapes es una lista -> declaramos lista con None
+                "img_shapes": [None],
             }
-            # Solo agregamos las que existan
+
+            # Aplicar solo los keys existentes
             for k, v in static_shapes.items():
-                if k in dynamic_shapes:
+                if k in call.kwargs:
                     dynamic_shapes[k] = v
 
             print("🚀 Exportando modelo con torch.export...")