up

app.py

@@ -11,15 +11,14 @@ dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the model pipeline
-pipe = DiffusionPipeline.from_pretrained(
-    "black-forest-labs/Flux.1-Dev", torch_dtype=dtype
-).to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", torch_dtype=dtype).to(device)
+
 
 @spaces.GPU(duration=120)
-def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken):
+def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken, progress=gr.Progress(track_tqdm=True)):
     if not filename.endswith(".pt2"):
         raise NotImplementedError("The filename must end with a `.pt2` extension.")
-    
+
     # this will throw if token is invalid
     try:
         _ = whoami(oauth_token.token)
@@ -27,12 +26,9 @@ def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken):
         # --- Ahead-of-time compilation ---
         compiled_transformer = compile_transformer(pipe, prompt="prompt")
 
-        token = oauth_token.token 
+        token = oauth_token.token
         out = _push_compiled_graph_to_hub(
-            compiled_transformer.archive_file,
-            repo_id=repo_id,
-            token=token,
-            path_in_repo=filename
+            compiled_transformer.archive_file, repo_id=repo_id, token=token, path_in_repo=filename
         )
         if not isinstance(out, str) and hasattr(out, "commit_url"):
             commit_url = out.commit_url
@@ -40,9 +36,12 @@ def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken):
         else:
             return out
     except Exception as e:
-        raise gr.Error(f"""Oops, you forgot to login. Please use the loggin button on the top left to migrate your repo {e}""")
-     
-css="""
+        raise gr.Error(
+            f"""Oops, you forgot to login. Please use the loggin button on the top left to migrate your repo {e}"""
+        )
+
+
+css = """
 #col-container {
     margin: 0 auto;
     max-width: 520px;
@@ -50,8 +49,12 @@ css="""
 """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("## Compile [Flux.1-Dev](https://hf.co/black-forest-labs/Flux.1-Dev) graph ahead of time & push to the Hub")
-        gr.Markdown("Enter a **repo_id** and **filename**. This repo automatically compiles the Flux.1-Dev model ahead of time. Read more about this in [this post](https://huggingface.co/blog/zerogpu-aoti).")
+        gr.Markdown(
+            "## Compile [Flux.1-Dev](https://hf.co/black-forest-labs/Flux.1-Dev) graph ahead of time & push to the Hub"
+        )
+        gr.Markdown(
+            "Enter a **repo_id** and **filename**. This repo automatically compiles the Flux.1-Dev model ahead of time. Read more about this in [this post](https://huggingface.co/blog/zerogpu-aoti)."
+        )
 
         repo_id = gr.Textbox(label="repo_id", placeholder="e.g. sayakpaul/qwen-aot")
         filename = gr.Textbox(label="filename", placeholder="e.g. compiled.pt2")
@@ -62,17 +65,19 @@ with gr.Blocks(css=css) as demo:
 
     run.click(push_to_hub, inputs=[repo_id, filename], outputs=[markdown_out])
 
+
 def swap_visibilty(profile: gr.OAuthProfile | None):
     return gr.update(elem_classes=["main_ui_logged_in"]) if profile else gr.update(elem_classes=["main_ui_logged_out"])
-        
-css_login = '''
+
+
+css_login = """
 .main_ui_logged_out{opacity: 0.3; pointer-events: none; margin: 0 auto; max-width: 520px}
-'''
+"""
 with gr.Blocks(css=css_login) as demo_login:
     gr.LoginButton()
     with gr.Column(elem_classes="main_ui_logged_out") as main_ui:
         demo.render()
     demo_login.load(fn=swap_visibilty, outputs=main_ui)
-    
+
 demo_login.queue()
-demo_login.launch()
+demo_login.launch()

hub_utils.py

@@ -1,10 +1,11 @@
 from io import BytesIO
 from huggingface_hub import create_repo, upload_file
-import tempfile 
+import tempfile
 import os
 
 DEFAULT_ARCHIVE_FILENAME = "archived_graph.pt2"
 
+
 def _push_compiled_graph_to_hub(archive: BytesIO, repo_id, **kwargs):
     if not isinstance(archive, BytesIO):
         raise NotImplementedError("Incorrect type of `archive` provided.")
@@ -13,9 +14,7 @@ def _push_compiled_graph_to_hub(archive: BytesIO, repo_id, **kwargs):
     private = kwargs.pop("private", False)
     path_in_repo = kwargs.pop("path_in_repo", DEFAULT_ARCHIVE_FILENAME)
     token = kwargs.pop("token")
-    repo_id = create_repo(
-        repo_id, private=private, exist_ok=True, token=token
-    ).repo_id
+    repo_id = create_repo(repo_id, private=private, exist_ok=True, token=token).repo_id
 
     with tempfile.TemporaryDirectory() as tmpdir:
         output_path = os.path.join(tmpdir, os.path.basename(path_in_repo))
@@ -24,8 +23,8 @@ def _push_compiled_graph_to_hub(archive: BytesIO, repo_id, **kwargs):
 
         try:
             info = upload_file(
-                repo_id=repo_id, 
-                path_or_fileobj=output_path, 
+                repo_id=repo_id,
+                path_or_fileobj=output_path,
                 path_in_repo=os.path.basename(path_in_repo),
                 commit_message=commit_message,
                 token=token,
@@ -33,4 +32,4 @@ def _push_compiled_graph_to_hub(archive: BytesIO, repo_id, **kwargs):
             return info
         except Exception as e:
             print(f"File couldn't be pushed to the Hub with the following error: {e}.")
-            return e
+            return e

optimization.py

@@ -5,26 +5,27 @@ import spaces
 import torch
 from torch.utils._pytree import tree_map
 
-P = ParamSpec('P')
+P = ParamSpec("P")
 
-TRANSFORMER_HIDDEN_DIM = torch.export.Dim('hidden', min=4096, max=8212)
+TRANSFORMER_HIDDEN_DIM = torch.export.Dim("hidden", min=4096, max=8212)
 
 # Specific to Flux. More about this is available in
 # https://huggingface.co/blog/zerogpu-aoti
 TRANSFORMER_DYNAMIC_SHAPES = {
-    'hidden_states': {1: TRANSFORMER_HIDDEN_DIM},
-    'img_ids': {0: TRANSFORMER_HIDDEN_DIM},
+    "hidden_states": {1: TRANSFORMER_HIDDEN_DIM},
+    "img_ids": {0: TRANSFORMER_HIDDEN_DIM},
 }
 
 INDUCTOR_CONFIGS = {
-    'conv_1x1_as_mm': True,
-    'epilogue_fusion': False,
-    'coordinate_descent_tuning': True,
-    'coordinate_descent_check_all_directions': True,
-    'max_autotune': True,
-    'triton.cudagraphs': True,
+    "conv_1x1_as_mm": True,
+    "epilogue_fusion": False,
+    "coordinate_descent_tuning": True,
+    "coordinate_descent_check_all_directions": True,
+    "max_autotune": True,
+    "triton.cudagraphs": True,
 }
 
+
 def compile_transformer(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
     @spaces.GPU(duration=1500)
     def f():
@@ -35,12 +36,9 @@ def compile_transformer(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.k
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
 
         exported = torch.export.export(
-            mod=pipeline.transformer, 
-            args=call.args, 
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes
+            mod=pipeline.transformer, args=call.args, kwargs=call.kwargs, dynamic_shapes=dynamic_shapes
         )
         return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
-    
+
     compiled_transformer = f()
-    return compiled_transformer
+    return compiled_transformer