#!/usr/bin/env python

import gradio as gr
import PIL.Image
import os
from gradio_client import Client

lgm_mini_client = Client("dylanebert/LGM-mini")
triposr_client = Client("stabilityai/TripoSR")

def run(image, model_name):
    file_path = "temp.png"
    image.save(file_path)
    if model_name=='lgm-mini':
      result = lgm_mini_client.predict(
      file_path,	# filepath  in 'image' Image component
      api_name="/run"
      )
      output = result
    elif model_name=='triposr':
        
        process_result = triposr_client.predict(
      file_path,	# filepath  in 'Input Image' Image component
      True,	# bool  in 'Remove Background' Checkbox component
      0.85,	# float (numeric value between 0.5 and 1.0) in 'Foreground Ratio' Slider component
      api_name="/preprocess")
        
        result = triposr_client.predict(
		process_result,	# filepath  in 'Processed Image' Image component
		256,	# float (numeric value between 32 and 320) in 'Marching Cubes Resolution' Slider component
		api_name="/generate")

        output = result[0]
    return output


demo = gr.Interface(
    fn=run,
    inputs=[gr.Image(type="pil"),gr.Textbox(label="Model Name")],
    outputs=gr.Model3D(label="3D Model"),
    api_name="/synthesize"
)

demo.launch()