cbensimon

"""
"""

import spaces
import torch
from diffusers import DiffusionPipeline
from diffusers.utils import load_image
from spaces.zero.torch.aoti import ZeroGPUCompiledModel
from spaces.zero.torch.aoti import ZeroGPUWeights
from torch.utils._pytree import tree_map

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)