File size: 1,795 Bytes
fbf185f 0fa316d ab4ccf4 0fa316d ab4ccf4 0fa316d fbf185f 0fa316d fbf185f 06b8314 52a6a99 fbf185f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
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
|