import gradio as gr
from huggingface_hub import hf_hub_download
from cadfusion.inference import InferenceRunner
from cadfusion.utils.config import load_config

# Download config & model from HF Hub
CONFIG_PATH = hf_hub_download(
    repo_id="microsoft/CADFusion",
    filename="adapter_config.json",
    revision="v1_1"
)

CHECKPOINT_PATH = hf_hub_download(
    repo_id="microsoft/CADFusion",
    filename="adapter_model.safetensors",
    revision="v1_1"
)

# Initialize model
config = load_config(CONFIG_PATH)
runner = InferenceRunner(config=config)
runner.load_checkpoint(CHECKPOINT_PATH)

# Inference function
def generate_cad(prompt: str):
    outputs = runner.infer_from_text(prompt)
    out_file = "output.stl"
    mesh = outputs.to_trimesh()
    mesh.export(out_file)
    return out_file

demo = gr.Interface(
    fn=generate_cad,
    inputs=gr.Textbox(label="Enter CAD description"),
    outputs=gr.File(label="Generated STL Model"),
    title="CADFusion - Text to CAD",
    description="Enter a text prompt to generate a 3D STL mesh."
)

if __name__ == "__main__":
    demo.launch()