import os
import math
import numpy as np
import onnxruntime as ort
from PIL import Image

LATENT_FEATURES = 512
RESOLUTION = 256
LAST_INDEX = math.log2(RESOLUTION) - 2

MODEL_PATH = os.path.join("model", "batik_stylegan.onnx")
model = ort.InferenceSession(MODEL_PATH)

alpha = np.array([1.0], dtype=np.float32)
steps = np.array([LAST_INDEX], dtype=np.int64)

def generate_stylegan():
    z = np.random.randn(1, LATENT_FEATURES).astype(np.float32)

    output = model.run(None, {
        'z': z,
        'alpha': alpha,
        'steps': steps
    })[0]

    image = output.squeeze(0)
    image = (image * 0.5 + 0.5) * 255
    image = image.astype(np.uint8)
    image = np.transpose(image, (1, 2, 0))
    pil_img = Image.fromarray(image, 'RGB')
    
    return pil_img.resize((512, 512), Image.LANCZOS)