u[

optimization.py

@@ -100,14 +100,15 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
         compiled_landscape = aoti_compile(exported_landscape, INDUCTOR_CONFIGS)
         compiled_portrait = aoti_compile(exported_portrait, INDUCTOR_CONFIGS)
-        # compiled_portrait.weights = (
-        #     compiled_landscape.weights
-        # )  # Avoid weights duplication when serializing back to main process
+        compiled_portrait.weights = (
+            compiled_landscape.weights
+        )  # Avoid weights duplication when serializing back to main process
 
         return compiled_landscape, compiled_portrait
 
     compiled_landscape, compiled_portrait = compile_transformer()
 
+    @torch.no_grad()
     def combined_transformer(*args, **kwargs):
         hidden_states: torch.Tensor = kwargs["hidden_states"]
         unpacked_hidden_states = LTXConditionPipeline._unpack_latents(

optimization_utils.py

@@ -67,7 +67,7 @@ def aoti_compile(
     files: list[str | Weights] = [file for file in artifacts if isinstance(file, str)]
     package_aoti(archive_file, files)
     (weights,) = (artifact for artifact in artifacts if isinstance(artifact, Weights))
-    zerogpu_weights = ZeroGPUWeights({name: weights.get_weight(name)[0] for name in weights})
+    zerogpu_weights = ZeroGPUWeights({name: weights.get_weight(name)[0] for name in weights}, to_cuda=True)
     return ZeroGPUCompiledModel(archive_file, zerogpu_weights)