app.py

import gradio as gr
import spaces
import torch
from diffusers import QwenImagePipeline
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
from optimization import compile_transformer
from hub_utils import _push_compiled_graph_to_hub
from huggingface_hub import whoami

# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the model pipeline
pipe = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=dtype).to(device)
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())

@spaces.GPU(duration=120)
def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken):
    if not filename.endswith(".pt2"):
        raise NotImplementedError("The filename must end with a `.pt2` extension.")
    
    # this will throw if token is invalid
    _ = whoami(oauth_token.token)

    # --- Ahead-of-time compilation ---
    compiled_transformer = compile_transformer(pipe, prompt="prompt")

    token = oauth_token.token 
    out = _push_compiled_graph_to_hub(
        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
        return f"[{commit_url}]({commit_url})"
    else:
        return out
     

css = """
#app {max-width: 840px; margin: 0 auto;}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("## Compile a model graph ahead of time & push to the Hub")
    gr.Markdown("Enter a **repo_id** and **filename**. This repo automatically compiles the [QwenImage](https://hf.co/Qwen/Qwen-Image) model on start.")

    with gr.Row():
        repo_id = gr.Textbox(label="repo_id", placeholder="e.g. sayakpaul/qwen-aot")
        filename = gr.Textbox(label="filename", placeholder="e.g. compiled.pt2")

    run = gr.Button("Push graph to Hub", variant="primary")

    markdown_out = gr.Markdown()

    run.click(push_to_hub, inputs=[repo_id, filename], outputs=[markdown_out])

if __name__ == "__main__":
    demo.launch(show_error=True)