Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,652 Bytes
835cf28 554e4ad 835cf28 554e4ad 835cf28 554e4ad 835cf28 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import spaces
import torch
from typing import Callable, ParamSpec, Any
from spaces.zero.torch.aoti import ZeroGPUCompiledModel, ZeroGPUWeights
P = ParamSpec('P')
# Z-Image sequences vary based on resolution.
TRANSFORMER_SEQ_LENGTH = torch.export.Dim('seq_length', min=256, max=16384)
# forward(self, x, attn_mask, freqs_cis, adaln_input)
TRANSFORMER_DYNAMIC_SHAPES = (
{1: TRANSFORMER_SEQ_LENGTH}, # x
{1: TRANSFORMER_SEQ_LENGTH}, # attn_mask
{1: TRANSFORMER_SEQ_LENGTH}, # freqs_cis
None, # adaln_input
)
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,
}
def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
target_layers = pipeline.transformer.layers
@spaces.GPU(duration=1500)
def compile_block():
block = target_layers[0]
with spaces.aoti_capture(block) as call:
pipeline(*args, **kwargs)
dynamic_shapes = TRANSFORMER_DYNAMIC_SHAPES
with torch.no_grad():
exported = torch.export.export(
mod=block,
args=call.args,
kwargs=call.kwargs,
dynamic_shapes=dynamic_shapes,
)
return spaces.aoti_compile(exported, INDUCTOR_CONFIGS).archive_file
archive_file = compile_block()
for block in target_layers:
weights = ZeroGPUWeights(block.state_dict())
block.forward = ZeroGPUCompiledModel(archive_file, weights) |