from asyncio import constants
import gradio as gr
import requests
import os 
from base64 import b64decode
from PIL import Image
import io
import numpy as np


def generate_image(seed,psi):
  iface = gr.Interface.load("spaces/hysts/StyleGAN-Human")
  print("about to die",iface,dir(iface))
  print(iface.blocks,iface.children,iface.fns)
  img=iface.fns[0].fn(seed,psi)
  
  #wrong format, gah! convert to numpy array
  header, encoded = img.split(",", 1)
  data = b64decode(encoded)
  image = Image.open(io.BytesIO(data))
  image_np = np.array(image)
  
  print('got img',img)

  return image_np

def generate_model(img):
    iface = gr.Interface.load("spaces/nagolinc/styleGanHuman_and_PIFu")
    model,file=iface(img)
    return model,file


demo = gr.Blocks()

with demo:
  gr.Markdown("<h1><center>StyleGan-Human + PIFu </center></h1>")
  gr.Markdown(
        "create a person and then generate a model from that person's image"
    )
    
  
  with gr.Row():
    b0 = gr.Button("generate image")
    b1 = gr.Button("generate model")
  
  with gr.Row():  
    seed=gr.Number(default=0, label='Seed')
    psi=gr.Slider(0, 2, step=0.05, default=0.7, label='Truncation psi')
    #outputImage = gr.Image(label="portrait",type="filepath", shape=(256,256))
    output_image = gr.outputs.Image(type='numpy', label='Output')
    model = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model")
    file= gr.File(label="Download 3D Model")

  
  b0.click(generate_image,inputs=[seed,psi],outputs=output_image)
  b1.click(generate_model, inputs=output_image, outputs=[model,file])

demo.launch(enable_queue=True, debug=True)