Spaces:
Sleeping
Sleeping
File size: 1,108 Bytes
d16b52d |
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 |
import torch
import torch.nn as nn
from diffusers import AutoencoderKL
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import (
retrieve_latents,
)
from .builder import EngineBuilder
from .models import BaseModel
class TorchVAEEncoder(torch.nn.Module):
def __init__(self, vae: AutoencoderKL):
super().__init__()
self.vae = vae
def forward(self, x: torch.Tensor):
return retrieve_latents(self.vae.encode(x))
def compile_engine(
torch_model: nn.Module,
model_data: BaseModel,
onnx_path: str,
onnx_opt_path: str,
engine_path: str,
opt_image_height: int = 512,
opt_image_width: int = 512,
opt_batch_size: int = 1,
engine_build_options: dict = {},
):
builder = EngineBuilder(
model_data,
torch_model,
device=torch.device("cuda"),
)
builder.build(
onnx_path,
onnx_opt_path,
engine_path,
opt_image_height=opt_image_height,
opt_image_width=opt_image_width,
opt_batch_size=opt_batch_size,
**engine_build_options,
)
|