from cog import BasePredictor, Input, Path
import tempfile
import time
from meshgpt_pytorch import MeshTransformer, mesh_render
from os.path import dirname

class Predictor(BasePredictor):
    def setup(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.transformer = MeshTransformer.from_pretrained("MarcusLoren/MeshGPT-preview")

    def predict(
        self,
        text: str = Input(description="Enter labels, separated by commas"),
        num_input: int = Input(description="Number of examples per input", default=1),
        num_temp: float = Input(description="Temperature (0 to 1)", default=0),
    ) -> Path:
        """Run a single prediction on the model"""
        self.transformer.eval()
        labels = [label.strip() for label in text.split(",")]

        output = []
        current_time = time.time()

        formatted_time = time.strftime(
            "%Y-%m-%d %H:%M:%S", time.localtime(current_time)
        )
        print(
            formatted_time,
            " Input:",
            text,
            "num_input",
            num_input,
            "num_temp",
            num_temp,
        )

        if num_input > 1:
            for label in labels:
                output.append(
                    (
                        self.transformer.generate(
                            texts=[label] * num_input, temperature=num_temp
                        )
                    )
                )
        else:
            output.append(
                (self.transformer.generate(texts=labels, temperature=num_temp))
            )
        file_name = "./mesh.obj"
        file_path = Path(tempfile.mkdtemp()) / file_name
        mesh_render.save_rendering(str(file_path), output)
        return file_path