Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,203 Bytes
f50853b 7655b4c f50853b f876753 fc44d4b f876753 7655b4c f876753 7655b4c f876753 7655b4c f876753 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
try:
import spaces
except ImportError:
# Define a dummy decorator if spaces is not available
def GPU(func):
return func
spaces = type('spaces', (), {'GPU': GPU})
import os
import torch
import argparse
from typing import *
from diffusers import StableDiffusionPipeline
from collections import deque
from triplaneturbo_executable.utils.mesh_exporter import export_obj
from triplaneturbo_executable import TriplaneTurboTextTo3DPipeline, TriplaneTurboTextTo3DPipelineConfig
# Initialize configuration and parameters
prompt = "a beautiful girl"
output_dir = "output"
adapter_name_or_path = "pretrained/triplane_turbo_sd_v1.pth"
num_results_per_prompt = 1
seed = 42
device = "cuda"
max_obj_files = 100
# download pretrained models if not exist
if not os.path.exists(adapter_name_or_path):
print(f"Downloading pretrained models from huggingface")
os.system(
f"huggingface-cli download --resume-download ZhiyuanthePony/TriplaneTurbo \
--include \"triplane_turbo_sd_v1.pth\" \
--local-dir ./pretrained \
--local-dir-use-symlinks False"
)
# Initialize the TriplaneTurbo pipeline
triplane_turbo_pipeline = TriplaneTurboTextTo3DPipeline.from_pretrained(adapter_name_or_path)
triplane_turbo_pipeline.to(device)
@spaces.GPU
def generate_3d_model(prompt, num_results_per_prompt=1, seed=42, device="cuda"):
"""
Generate 3D models using TriplaneTurbo pipeline.
Args:
prompt (str): Text prompt for the 3D model
num_results_per_prompt (int): Number of results to generate
seed (int): Random seed for generation
device (str): Device to use for computation
Returns:
dict: Output from the pipeline
"""
output = triplane_turbo_pipeline(
prompt=prompt,
num_results_per_prompt=num_results_per_prompt,
generator=torch.Generator(device=device).manual_seed(seed),
device=device,
)
# Initialize a deque with maximum length of 100 to store obj file paths
obj_file_queue = deque(maxlen=max_obj_files)
# Save mesh
os.makedirs(output_dir, exist_ok=True)
for i, mesh in enumerate(output["mesh"]):
vertices = mesh.v_pos
# 1. First rotate -90 degrees around X-axis to make the model face up
vertices = torch.stack([
vertices[:, 0], # x remains unchanged
vertices[:, 2], # y = z
-vertices[:, 1] # z = -y
], dim=1)
# 2. Then rotate 90 degrees around Y-axis to make the model face the observer
vertices = torch.stack([
-vertices[:, 2], # x = -z
vertices[:, 1], # y remains unchanged
vertices[:, 0] # z = x
], dim=1)
mesh.v_pos = vertices
# If mesh has normals, they need to be rotated in the same way
if mesh.v_nrm is not None:
normals = mesh.v_nrm
# 1. Rotate -90 degrees around X-axis
normals = torch.stack([
normals[:, 0],
normals[:, 2],
-normals[:, 1]
], dim=1)
# 2. Rotate 90 degrees around Y-axis
normals = torch.stack([
-normals[:, 2],
normals[:, 1],
normals[:, 0]
], dim=1)
mesh._v_nrm = normals
# Save obj file and add its path to the queue
name = f"{prompt.replace(' ', '_')}_{seed}_{i}"
save_paths = export_obj(mesh, f"{output_dir}/{name}.obj")
obj_file_queue.append(save_paths[0])
# If an old file needs to be removed (queue is at max length)
# and the file exists, delete it
if len(obj_file_queue) == max_obj_files and os.path.exists(obj_file_queue[0]):
old_file = obj_file_queue[0]
try:
os.remove(old_file)
except OSError as e:
print(f"Error deleting file {old_file}: {e}")
# Run the pipeline
output = generate_3d_model(
prompt=prompt,
num_results_per_prompt=num_results_per_prompt,
seed=seed,
device=device
)
|