Move README template in private and pass it

job.py

@@ -126,40 +126,13 @@ def generate_samples(samples_dir: Path):
     output.images[0].save(f'{samples_dir}/edited.png')
 
 
-README_TEMPLATE = textwrap.dedent(
-    """
-    ## How to use
-    ``` python
-    {MODEL_INIT}\n
-    {AOTI_LOAD}
-    ```
-
-    ## How to reproduce or customize
-    ``` bash
-    pip install huggingface_hub
-
-    # Get the job file and edit if needed
-    hf download {REPO_ID} job.py
-
-    # Run the job
-    hf jobs uv run job.py \
-        --hardware {JOB_HARDWARE} # Change if needed \
-        --image {JOB_IMAGE} # ... \
-        -e PARENT_REPO_ID={REPO_ID} \
-        -e OUTPUT_REPO_ID=<output-repo-id> \
-        -e HF_TOKEN=<hf-token>
-    ```
-    """
-)
-
-
 # ========================================================
 # ===== Meta Job below this line. Do not edit unless =====
 # ========================================================
 
 
 def capture_model_init_region():
-    source = inspect.getsource(sys.modules[__name__])
+    source = inspect.getsource(sys.modules['__main__'])
     return (source
         .split('# README::MODEL_INIT::START')[1]
         .split('# README::MODEL_INIT::END')[0]
@@ -198,6 +171,7 @@ def create_aoti_repo(
     module_expr: str,
     compile_module: Callable[[torch.nn.Module, Path]],
     generate_samples: Callable[[Path]],
+    readme_template: str,
     aoti_loader: Callable[[torch.nn.Module, Path], None] | None = None,
 ):
     """
@@ -232,7 +206,7 @@ def create_aoti_repo(
             repo_id=output_repo_id,
             aoti_loader=aoti_loader,
         )
-        readme_contents = README_TEMPLATE.format(
+        readme_contents = readme_template.format(
             MODEL_INIT=capture_model_init_region(),
             AOTI_LOAD=aoti_load_readme,
             REPO_ID=output_repo_id,
@@ -266,6 +240,39 @@ def create_aoti_repo(
         upload_folder(repo_id=output_repo_id, folder_path=tempdir)
 
 
+README_TEMPLATE = textwrap.dedent(
+    """
+    ## How to use
+    ``` python
+    {MODEL_INIT}\n
+    {AOTI_LOAD}
+    ```
+
+    ## How to reproduce or customize
+    ``` bash
+    pip install huggingface_hub
+
+    # Get the job file and edit if needed
+    hf download {REPO_ID} job.py
+
+    # Run the job
+    hf jobs uv run job.py \
+        --hardware {JOB_HARDWARE} # Change if needed \
+        --image {JOB_IMAGE} # ... \
+        -e PARENT_REPO_ID={REPO_ID} \
+        -e OUTPUT_REPO_ID=<output-repo-id> \
+        -e HF_TOKEN=<hf-token>
+    ```
+
+    ## Environment
+    <details>
+    <summary>Click to expand</summary>
+    {ENVIRONMENT}
+    </details>
+    """
+)
+
+
 if __name__ == '__main__':
     create_aoti_repo(
         output_repo_id=os.environ['OUTPUT_REPO_ID'],
@@ -273,5 +280,6 @@ if __name__ == '__main__':
         module_expr='pipeline.transformer',
         compile_module=compile_module,
         generate_samples=generate_samples,
+        readme_template=README_TEMPLATE, # TODO: allow for extra template data dict?
         aoti_loader=aoti_loader,
     )