cbensimon's picture
cbensimon HF Staff
ZeroGPUCompiledModel
674e245
raw
history blame
2.69 kB
"""
"""
from datetime import datetime
# Upgrade PyTorch
import os
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 torch torchvision spaces')
# CUDA toolkit install
t0 = datetime.now()
from utils.cuda_toolkit import install_cuda_toolkit; install_cuda_toolkit()
print('install_cuda_toolkit', -(t0 - (t0 := datetime.now())))
# Actual app.py
import os
import gradio as gr
import spaces
import torch
import torch._inductor
from diffusers import FluxPipeline
from .utils.zerogpu import aoti_compile
pipeline = FluxPipeline.from_pretrained('black-forest-labs/FLUX.1-schnell', torch_dtype=torch.bfloat16).to('cuda')
print('FluxPipeline.from_pretrained', -(t0 - (t0 := datetime.now())))
package_path = 'pipeline.pt2'
@spaces.GPU(duration=1500)
def compile_transformer():
def _example_tensor(*shape):
return torch.randn(*shape, device='cuda', dtype=torch.bfloat16)
is_timestep_distilled = not pipeline.transformer.config.guidance_embeds
seq_length = 256 if is_timestep_distilled else 512
transformer_kwargs = {
'hidden_states': _example_tensor(1, 4096, 64),
'timestep': torch.tensor([1.], device='cuda', dtype=torch.bfloat16),
'guidance': None if is_timestep_distilled else torch.tensor([1.], device='cuda', dtype=torch.bfloat16),
'pooled_projections': _example_tensor(1, 768),
'encoder_hidden_states': _example_tensor(1, seq_length, 4096),
'txt_ids': _example_tensor(seq_length, 3),
'img_ids': _example_tensor(4096, 3),
'joint_attention_kwargs': {},
'return_dict': False,
}
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,
}
exported = torch.export.export(pipeline.transformer, args=(), kwargs=transformer_kwargs)
return aoti_compile(exported, inductor_configs)
transformer_config = pipeline.transformer.config
pipeline.transformer = compile_transformer()
pipeline.transformer.config = transformer_config
@spaces.GPU
def _generate_image(prompt: str, t0: datetime):
print('@spaces.GPU', -(t0 - (t0 := datetime.now())))
images = []
for _ in range(4):
images += pipeline(prompt, num_inference_steps=4).images
print('pipeline', -(t0 - (t0 := datetime.now())))
return images
def generate_image(prompt: str, progress=gr.Progress(track_tqdm=True)):
return _generate_image(prompt, datetime.now())
gr.Interface(generate_image, gr.Text(), gr.Gallery()).launch(show_error=True)