cbensimon

@spaces.GPU
def compile():
    with spaces.aoti_capture(pipe.transformer) as call_landscape:
        pipe("prompt", width=832, height=480)
    with spaces.aoti_capture(pipe.transformer) as call_portrait:
        pipe("prompt", width=480, height=832)

    exported_landscape = torch.export.export(model, args=call_landscape.args, kwargs=call_landscape.kwargs)
    exported_portrait = torch.export.export(model, args=call_portrait.args, kwargs=call_portrait.kwargs)