COMPILED_TRANSFORMER_1 = compiled_transformer_1

optimization.py

@@ -23,6 +23,10 @@ from optimization_utils import ZeroGPUCompiledModel, ZeroGPUWeights
 
 P = ParamSpec('P')
 
+# Expose compiled models so app.py can offer them for download
+COMPILED_TRANSFORMER_1 = None
+COMPILED_TRANSFORMER_2 = None
+
 LATENT_FRAMES_DIM = torch.export.Dim('num_latent_frames', min=8, max=81)
 LATENT_PATCHED_HEIGHT_DIM = torch.export.Dim('latent_patched_height', min=30, max=52)
 LATENT_PATCHED_WIDTH_DIM = torch.export.Dim('latent_patched_width', min=30, max=52)
@@ -99,6 +103,8 @@ def load_compiled_transformers_from_hub(
 
 
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
+    global COMPILED_TRANSFORMER_1, COMPILED_TRANSFORMER_2
+
     @spaces.GPU(duration=1500)
     def compile_transformer():
         pipeline.load_lora_weights(
@@ -156,6 +162,10 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
     else:
         compiled_transformer_1, compiled_transformer_2 = compile_transformer()
 
+    # expose for downloads
+    COMPILED_TRANSFORMER_1 = compiled_transformer_1
+    COMPILED_TRANSFORMER_2 = compiled_transformer_2
+
     pipeline.transformer.forward = compiled_transformer_1
     drain_module_parameters(pipeline.transformer)