#!/usr/bin/env python

import gradio as gr
import PIL.Image

from gradio_client import Client

shap_e_client = Client("hysts/Shap-E")
triposr_client = Client("stabilityai/TripoSR")


def run(image, model_name):
    file_path = "temp.png"
    image.save(file_path)
    print(file_path)
    model_name = model_name.lower()
    if model_name=='shap-e':
        result = shap_e_client.predict(
		file_path,	# filepath  in 'Input image' Image component
		0,	# float (numeric value between 0 and 2147483647) in 'Seed' Slider component
		1,	# float (numeric value between 1 and 20) in 'Guidance scale' Slider component
		2,	# float (numeric value between 2 and 100) in 'Number of inference steps' Slider component
		api_name="/image-to-3d")
        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.5,	# float (numeric value between 0.5 and 1.0) in 'Foreground Ratio' Slider component
		api_name="/preprocess")
        print(type(process_result))
        
        result = triposr_client.predict(
		process_result,	# filepath  in 'Processed Image' Image component
		32,	# float (numeric value between 32 and 320) in 'Marching Cubes Resolution' Slider component
		api_name="/generate")

        output = result[1]
            
    return output
        
    
with gr.Blocks() as demo:
    with gr.Group():
        image = gr.Image(label="Input image", show_label=False, type="pil")
        model_name = gr.Textbox(label="Model name", show_label=False)
        run_button = gr.Button("Run")
        result = gr.Model3D(label="Result", show_label=False)

    run_button.click(
        fn=run,
        inputs=[
            image,
            model_name
        ],
        outputs=result,
        api_name="synthesize"
    )

demo.launch()