Update optimization.py

optimization.py

@@ -1,7 +1,6 @@
 """
 Optimization module for Qwen-Image-Edit using TorchAO quantization and AoTI compilation.
 """
-
 from typing import Any
 from typing import Callable
 from typing import ParamSpec
@@ -47,21 +46,18 @@ INDUCTOR_CONFIGS = {
 
 
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
-    """
-    Optimizes the Qwen-Image-Edit pipeline using AoT compilation and quantization.
-    This function pre-compiles the transformer for faster inference.
-    """
+
     @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
 
-        # Optional: Uncomment to enable Float8 quantization
         # quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
-
+        
         exported = torch.export.export(
             mod=pipeline.transformer,
             args=call.args,
@@ -71,4 +67,4 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
         return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
 
-    spaces.aoti_apply(compile_transformer(), pipeline.transformer)
+    spaces.aoti_apply(compile_transformer(), pipeline.transformer)