Almost the right structure

job.py

@@ -9,13 +9,19 @@
 #   "setuptools",
 # ]
 # ///
+# fmt: off
 
+#region how-to
 import spaces
 import torch
 from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
-from PIL import Image
-from spaces.zero.torch.aoti import LazyAOTIModel
-from torch.utils._pytree import tree_map
+from huggingface_hub import hf_hub_download
+
+pipeline = Flux2KleinPipeline.from_pretrained(
+    'black-forest-labs/FLUX.2-klein-4B',
+    torch_dtype=torch.bfloat16
+).to('cuda')
+#endregion
 
 
 TRANSFORMER_IMAGE_DIM = torch.export.Dim('image_seq_length', min=4096, max=16384) # min: 0 images, max: 3 (1024x1024) images
@@ -51,64 +57,82 @@ INDUCTOR_CONFIGS = {
 }
 
 
-pipeline = Flux2KleinPipeline.from_pretrained(
-    'black-forest-labs/FLUX.2-klein-4B',
-    torch_dtype=torch.bfloat16
-).to('cuda')
+def compile_module(module: torch.nn.Module, output_dir: str):
+    from PIL import Image
+    from torch.utils._pytree import tree_map
 
+    blocks = {
+        'double': module.transformer_blocks,
+        'single': module.single_transformer_blocks,
+    }
 
-def generate_sample(filename: str):
-    pipeline(
-        prompt="Remove the sunglasses",
-        image=[Image.open('cat-sunglasses.jpg')],
-        guidance_scale=2.5,
-        generator=torch.Generator(device='cuda').manual_seed(42),
-    ).images[0].save(filename) # pyright: ignore[reportAttributeAccessIssue]
-
-
-generate_sample('before.png')
-
-
-blocks = {
-    'double': pipeline.transformer.transformer_blocks,
-    'single': pipeline.transformer.single_transformer_blocks,
-}
+    def compile_block(blocks_kind: str):
+        block = blocks[blocks_kind][0]
+        with spaces.aoti_capture(block) as call:
+            pipeline(
+                prompt="prompt",
+                image=[Image.new("RGB", (1024, 1024))],
+            )
 
-def compile_block(blocks_kind: str):
-    block = blocks[blocks_kind][0]
-    with spaces.aoti_capture(block) as call:
-        pipeline(
-            prompt="prompt",
-            image=[Image.new("RGB", (1024, 1024))],
-        )
+        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES[blocks_kind]
 
-    dynamic_shapes = tree_map(lambda t: None, call.kwargs)
-    dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES[blocks_kind]
+        with torch.no_grad():
+            exported = torch.export.export(
+                mod=block,
+                args=call.args,
+                kwargs=call.kwargs,
+                dynamic_shapes=dynamic_shapes,
+            )
 
-    with torch.no_grad():
-        exported = torch.export.export(
-            mod=block,
-            args=call.args,
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes,
-        )
+        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS).archive_file
 
-    return spaces.aoti_compile(exported, INDUCTOR_CONFIGS).archive_file
+    for blocks_kind in ('double', 'single'):
+        with open(f'{output_dir}/{blocks_kind}.pt2', 'wb') as f:
+            archive = compile_block(blocks_kind)
+            f.write(archive.getbuffer())
 
-for blocks_kind in ('double', 'single'):
-    archive_file = compile_block(blocks_kind)
-    with open(f'{blocks_kind}.pt2', 'wb') as f:
-        f.write(archive_file.getbuffer()) # pyright: ignore[reportAttributeAccessIssue]
 
+def load_compiled(module: torch.nn.Module, repo_id: str, revision: str | None = None):
+    single = hf_hub_download(repo_id, 'single.pt2', revision=revision)
+    double = hf_hub_download(repo_id, 'double.pt2', revision=revision)
+    for block in module.single_transformer_blocks:
+        spaces.aoti_apply_graph(single, block)
+    for block in module.transformer_blocks:
+        spaces.aoti_apply_graph(double, block)
 
-# TODO: Do it with some spaces.aoti_ helper (lower-level aoti_blocks_load)
-for blocks_kind in ('double', 'single'):
-    archive_file_path = f'{blocks_kind}.pt2'
-    for block in blocks[blocks_kind]:
-        block.forward = LazyAOTIModel(archive_file_path).with_weights(block.state_dict())
 
-
-generate_sample('after.png')
+def generate_samples(output_dir: str):
+    output = pipeline(
+        prompt="Remove the sunglasses",
+        image=[Image.open('cat-sunglasses.jpg')],
+        guidance_scale=2.5,
+        generator=torch.Generator(device='cuda').manual_seed(42),
+    )
+    output.images[0].save(f'{output_dir}/edited.png') # pyright: ignore[reportAttributeAccessIssue]
+
+
+aot_init(
+    module=pipeline.transformer,
+    compile_module=compile_module,
+    load_compiled=load_compiled, # Optional (uses spaces.aoti_load as default)
+    generate_samples=generate_samples,
+    readme_template="""
+## How to use
+``` python
+{HOW_TO}
+{LOAD_COMPILED}
+```
+
+# How to customize
+``` bash
+``` bash
+pip install huggingface_hub
+hf download {REPO_ID} job.py
+hf jobs uv run --hardware h200 job.py
+```
+"""
+)
 
 # Push structure to hub
 # Push script to hub