Spaces:
Sleeping
Sleeping
File size: 3,294 Bytes
db66840 70a1838 d843d07 70a1838 db66840 70a1838 db66840 70a1838 d843d07 db66840 d843d07 db66840 d843d07 db66840 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | from stylegan_model import StyleGAN
from vanillagan_model import VanillaGAN
import torch
from io import BytesIO
from torchvision.utils import save_image
import numpy as np
import legacy
from PIL import Image
import time
import onnxruntime as ort
LATENT_FEATURES = 512
RESOLUTION = 128
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def load_model_pt(path='model_128.pt',model_type='stylegan'):
if model_type == "stylegan":
model = StyleGAN(LATENT_FEATURES, RESOLUTION).to(DEVICE)
last_checkpoint = torch.load(path, map_location=DEVICE)
model.load_state_dict(last_checkpoint['generator'], strict=False)
elif model_type == "vanillagan":
model = VanillaGAN(RESOLUTION, LATENT_FEATURES).to(DEVICE)
model.load_state_dict(torch.load(path, map_location=DEVICE))
model.eval()
return model
def generate_image_stylegan(generator, steps=5, alpha=1.0):
with torch.no_grad():
image = generator(torch.randn(1, LATENT_FEATURES, device=DEVICE), alpha=1.0, steps=steps)
image = image.tanh()
image = (image + 1) / 2
buffer = BytesIO()
save_image(image, buffer, format='PNG')
buffer.seek(0)
return buffer
def generate_image_vanillagan(generator):
with torch.no_grad():
image = generator(torch.randn(1, LATENT_FEATURES, device=DEVICE)).view(1, 3, RESOLUTION, RESOLUTION)
image = (image * 0.5 + 0.5).clamp(0, 1)
buffer = BytesIO()
save_image(image, buffer, format='PNG')
buffer.seek(0)
return buffer
def load_model_pkl(path='styleganv2.pkl'):
with open(path, 'rb') as f:
G = legacy.load_network_pkl(f)['G_ema'].to(DEVICE)
G.eval()
return G
def generate_image_from_pkl(generator, seed=0, trunc=1):
start = time.time()
z = torch.from_numpy(np.random.RandomState(seed).randn(1, generator.z_dim)).to(DEVICE)
label = torch.zeros(1, generator.c_dim, device=DEVICE)
img = generator(z, label, truncation_psi=trunc, noise_mode='const')
img = (img + 1) * (255 / 2)
img = img.clamp(0, 255).to(torch.uint8)
img = img[0].permute(1, 2, 0).cpu().numpy() # (Channel, Height, Width) to (Height, Width, Channel)
pil_image = Image.fromarray(img)
buffer = BytesIO()
pil_image.save(buffer, format='PNG')
buffer.seek(0)
end = time.time()
print(f"Image generation time: {end - start:.2f} seconds")
return buffer
def generate_image_from_onnx(path='model_128.onnx', model=None):
if model is None:
return ValueError("Model not provided.")
if model == 'progan' or model== 'dcgan':
z = np.random.randn(1, 512, 1, 1).astype(np.float32)
else:
z = np.random.randn(1, 512).astype(np.float32)
inference_session = ort.InferenceSession(path)
input_name = inference_session.get_inputs()[0].name
image = inference_session.run(None, {input_name: z})[0]
image = image.squeeze(0)
if model == "vanillagan":
image = image.reshape(3, 128, 128)
image = (image * 0.5 + 0.5) * 255
image = image.astype(np.uint8)
image = np.transpose(image, (1, 2, 0))
image = Image.fromarray(image, 'RGB')
buffer = BytesIO()
image.save(buffer, format='PNG')
buffer.seek(0)
return buffer
|