import gradio as gr
from diffusers import DiffusionPipeline
import torch
from gradio_client import Client
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
if torch.cuda.is_available(): pipe.to("cuda")
client = Client("stabilityai/triposr")


ABOUT_TEXT = """
# Text-to-3D with TripoSR + SDXL

Commercially-viable text-to-3D model. Usage must comply with the SDXL license.

For image-to-3D, use [TripoSR](https://huggingface.co/spaces/stabilityai/TripoSR) directly.
""".strip()
def generate(text):
    # generate image
    image = pipe(prompt=text).images[0]
    # preprocess
    result = client.predict(
		image,
		True,
		0.85,
        api_name="/preprocess"
    )
    result = client.predict(
		result,
        api_name="/generate"
    )
    return result
    
with gr.Blocks() as demo:
    gr.Markdown(ABOUT_TEXT)
    txt = gr.Textbox(interactive=True, label="Text instruction")
    btn = gr.Button("Generate")
    out = gr.Model3D(
        label="3D model",
        interactive=False,
    )
    btn.click(generate, inputs=txt, outputs=out)
demo.queue(api_open=False, default_concurrency_limit=20).launch(show_api=False)