import os
import sys

import gradio as gr
import numpy as np
import torch
from PIL import Image

sys.path.append(os.path.join(os.path.dirname(__file__), "stylegan2"))

import dnnlib
import legacy

DEVICE = "cpu"  # or "cuda"

device = torch.device(DEVICE)
with dnnlib.util.open_url("troublinggan-square.pkl") as f:
    G_square = legacy.load_network_pkl(f)["G_ema"].to(device)
with dnnlib.util.open_url("troublinggan-rectangle.pkl") as f:
    G_rectangle = legacy.load_network_pkl(f)["G_ema"].to(device)


def generate(sigma, model):
    G = {"v1": G_square, "v2": G_rectangle}[model]
    label = torch.zeros([1, G.c_dim], device=device)
    z = torch.from_numpy(np.random.RandomState().randn(1, G.z_dim) * sigma).to(device)
    img = G(
        z,
        label,
        truncation_psi=1,
        noise_mode="random",
        force_fp32=(DEVICE == "cpu"),
    )
    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
    img = Image.fromarray(img[0].cpu().numpy(), "RGB")
    if model == "v2":
        img = img.crop((0, 172, 1024, 1024 - 172))
    return img


iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Slider(
            value=1.0,
            minimum=0.5,
            maximum=2.0,
            step=0.25,
            info="perceived randomness of the image",
        ),
        gr.Radio(
            ["v1", "v2"],
            value="v1",
            info="v1: model containing photos from 2020 (resolution 512x512); v2: model containing photos from 2020-2021 (resolution 1024x680)",
        ),
    ],
    outputs=["image"],
    title="TroublingGAN",
    description="Generate a random image from the TroublingGAN model",
    article="""
TroublingGAN is a critical generative artwork and artistic research project that experiments with the StyleGAN vector-image generative neural network. It experiments with alternative use of GANs as a tool for artistic research.

Read more about TroublingGAN at [http://hamosova.com/TroublingGAN](http://hamosova.com/TroublingGAN)

_This research was created at the Academy of Performing Arts in Prague as part of the project "Extending the creative tools of Machine Learning and Artificial Intelligence - Experimental Tools in Artistic Practice" supported by the Ministry of Education and Science for specific university research at the Academy of Performing Arts in Prague in 2021._
""",
    allow_flagging="never",
)
iface.launch()