riciii7/FastAPI-Batik-GAN

1from stylegan_model import StyleGAN2from vanillagan_model import VanillaGAN3import torch4from io import BytesIO5from torchvision.utils import save_image6import numpy as np7import legacy8from PIL import Image9import time10import onnxruntime as ort11 12LATENT_FEATURES = 51213RESOLUTION = 12814 15DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')16def load_model_pt(path='model_128.pt',model_type='stylegan'):17    if model_type == "stylegan":18        model = StyleGAN(LATENT_FEATURES, RESOLUTION).to(DEVICE)19        last_checkpoint = torch.load(path, map_location=DEVICE)20        model.load_state_dict(last_checkpoint['generator'], strict=False)21    elif model_type == "vanillagan":22        model = VanillaGAN(RESOLUTION, LATENT_FEATURES).to(DEVICE)23        model.load_state_dict(torch.load(path, map_location=DEVICE))24    model.eval()25    return model26 27def generate_image_stylegan(generator, steps=5, alpha=1.0):28    with torch.no_grad():29        image = generator(torch.randn(1, LATENT_FEATURES, device=DEVICE), alpha=1.0, steps=steps)30        image = image.tanh()31        image = (image + 1) / 2 32 33        buffer = BytesIO()34        save_image(image, buffer, format='PNG')35        buffer.seek(0)36        return buffer37 38def generate_image_vanillagan(generator):39    with torch.no_grad():40        image = generator(torch.randn(1, LATENT_FEATURES, device=DEVICE)).view(1, 3, RESOLUTION, RESOLUTION)41        image = (image * 0.5 + 0.5).clamp(0, 1) 42 43        buffer = BytesIO()44        save_image(image, buffer, format='PNG')45        buffer.seek(0)46        return buffer47    48def load_model_pkl(path='styleganv2.pkl'):49    with open(path, 'rb') as f:50        G = legacy.load_network_pkl(f)['G_ema'].to(DEVICE)51    G.eval()52    return G53 54def generate_image_from_pkl(generator, seed=0, trunc=1):55    start = time.time()56    z = torch.from_numpy(np.random.RandomState(seed).randn(1, generator.z_dim)).to(DEVICE)57    label = torch.zeros(1, generator.c_dim, device=DEVICE)58    img = generator(z, label, truncation_psi=trunc, noise_mode='const')59    img = (img + 1) * (255 / 2)60    img = img.clamp(0, 255).to(torch.uint8)61    img = img[0].permute(1, 2, 0).cpu().numpy() # (Channel, Height, Width) to (Height, Width, Channel)62    pil_image = Image.fromarray(img)63 64    buffer = BytesIO()65    pil_image.save(buffer, format='PNG')66    buffer.seek(0)67 68    end = time.time()69    print(f"Image generation time: {end - start:.2f} seconds")70 71    return buffer72 73def generate_image_from_onnx(path='model_128.onnx', model=None):74    if model is None: 75        return ValueError("Model not provided.")76    if model == 'progan' or model== 'dcgan':77        z = np.random.randn(1, 512, 1, 1).astype(np.float32)78    else:79        z = np.random.randn(1, 512).astype(np.float32)80    inference_session = ort.InferenceSession(path)81    input_name = inference_session.get_inputs()[0].name82 83    image = inference_session.run(None, {input_name: z})[0]84 85    image = image.squeeze(0)86    87    if model == "vanillagan":88        image = image.reshape(3, 128, 128)89    image = (image * 0.5 + 0.5) * 25590    image = image.astype(np.uint8)91    image = np.transpose(image, (1, 2, 0))92    image = Image.fromarray(image, 'RGB')93 94    buffer = BytesIO()95    image.save(buffer, format='PNG')96    buffer.seek(0)97 98    return buffer99