import os
from gradio_client import Client, handle_file

class GlbController:
    def get_glb(self, hf_token='', base_image_url='', ss_guidance_strength=7.5,
                ss_sampling_steps=12, slat_guidance_strength=3,
                slat_sampling_steps=12, multiimage_algo='stochastic',
                mesh_simplify=0.95, texture_size=1024) -> str:
        if hf_token == '':
            hf_token = os.getenv("hf_token")

        client = Client("JeffreyXiang/TRELLIS", hf_token)

        client.predict(
            api_name="/start_session"
        )

        preprocess_image_result = client.predict(
            image=handle_file(base_image_url),
            api_name="/preprocess_image"
        )

        get_seed_result = client.predict(
            randomize_seed=True,
            seed=0,
            api_name="/get_seed"
        )

        client.predict(
            image=handle_file(preprocess_image_result),
            multiimages=[],
            seed=get_seed_result,
            ss_guidance_strength=ss_guidance_strength,
            ss_sampling_steps=ss_sampling_steps,
            slat_guidance_strength=slat_guidance_strength,
            slat_sampling_steps=slat_sampling_steps,
            multiimage_algo=multiimage_algo,
            api_name="/image_to_3d"
        )

        extract_glb = client.predict(
            mesh_simplify=mesh_simplify,
            texture_size=texture_size,
            api_name="/extract_glb"
        )

        return extract_glb