Spaces:
Running
on
Zero
Running
on
Zero
import os | |
import subprocess | |
import sys | |
try: | |
import spaces | |
except ImportError: | |
# Define a dummy decorator if spaces is not available | |
def GPU(func): | |
return func | |
spaces = type('spaces', (), {'GPU': GPU}) | |
# Check if setup has been run | |
setup_marker = ".setup_complete" | |
if not os.path.exists(setup_marker): | |
print("First run detected, installing dependencies...") | |
try: | |
subprocess.check_call(["bash", "setup.sh"]) | |
# Create marker file to indicate setup is complete | |
with open(setup_marker, "w") as f: | |
f.write("Setup completed") | |
print("Setup completed successfully!") | |
except subprocess.CalledProcessError as e: | |
print(f"Setup failed with error: {e}") | |
sys.exit(1) | |
import torch | |
import gradio as gr | |
from typing import * | |
from collections import deque | |
from diffusers import StableDiffusionPipeline | |
from triplaneturbo_executable import TriplaneTurboTextTo3DPipeline | |
from triplaneturbo_executable.utils.mesh_exporter import export_obj | |
# Initialize global variables | |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
ADAPTER_PATH = "pretrained/triplane_turbo_sd_v1.pth" #"/home/user/app/pretrained/triplane_turbo_sd_v1.pth" | |
PIPELINE = None # Will hold our pipeline instance | |
OBJ_FILE_QUEUE = deque(maxlen=100) # Queue to store OBJ file paths | |
def download_model(): | |
"""Download the pretrained model if not exists""" | |
if not os.path.exists(ADAPTER_PATH): | |
print("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" | |
) | |
def initialize_pipeline(): | |
"""Initialize the pipeline once and keep it in memory""" | |
global PIPELINE | |
if PIPELINE is None: | |
print("Initializing pipeline...") | |
PIPELINE = TriplaneTurboTextTo3DPipeline.from_pretrained(ADAPTER_PATH) | |
PIPELINE.to(DEVICE) | |
print("Pipeline initialized!") | |
return PIPELINE | |
def generate_3d_mesh(prompt: str) -> Tuple[str, str]: | |
"""Generate 3D mesh from text prompt""" | |
global PIPELINE, OBJ_FILE_QUEUE | |
# Use the global pipeline instance | |
pipeline = initialize_pipeline() | |
# Use fixed seed value | |
seed = 42 | |
# Generate mesh | |
output = pipeline( | |
prompt=prompt, | |
num_results_per_prompt=1, | |
generator=torch.Generator(device=DEVICE).manual_seed(seed), | |
) | |
# Save mesh | |
output_dir = "outputs" | |
os.makedirs(output_dir, exist_ok=True) | |
mesh_path = None | |
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 | |
name = f"{prompt.replace(' ', '_')}" | |
save_paths = export_obj(mesh, f"{output_dir}/{name}.obj") | |
mesh_path = save_paths[0] | |
# Add new file path to queue | |
OBJ_FILE_QUEUE.append(mesh_path) | |
# If queue is at max length, remove oldest file | |
if len(OBJ_FILE_QUEUE) == OBJ_FILE_QUEUE.maxlen: | |
old_file = OBJ_FILE_QUEUE[0] # Get oldest file (will be automatically removed from queue) | |
if os.path.exists(old_file): | |
try: | |
os.remove(old_file) | |
except OSError as e: | |
print(f"Error deleting file {old_file}: {e}") | |
return mesh_path, mesh_path # Return the path twice - once for 3D preview, once for download | |
with gr.Blocks(css=".output-image, .input-image, .image-preview {height: 512px !important}") as demo: | |
# Download model if needed | |
download_model() | |
# Initialize pipeline at startup | |
initialize_pipeline() | |
gr.Markdown( | |
""" | |
# 🌟 Text to 3D Mesh Generation with TriplaneTurbo | |
Demo of the paper "Progressive Rendering Distillation: Adapting Stable Diffusion for Instant Text-to-Mesh Generation without 3D Training Data" [CVPR 2025] | |
[GitHub Repository](https://github.com/theEricMa/TriplaneTurbo) | |
## Instructions | |
1. Enter a text prompt describing what 3D object you want to generate | |
2. Click "Generate" and wait for the model to create your 3D mesh | |
3. View the result in the 3D viewer or download the OBJ file | |
""" | |
) | |
with gr.Row(): | |
with gr.Column(scale=1): | |
prompt = gr.Textbox( | |
label="Text Prompt", | |
placeholder="Enter your text description...", | |
value="Armor dress style of outsiderzone fantasy helmet", | |
lines=2 | |
) | |
generate_btn = gr.Button("Generate", variant="primary") | |
examples = gr.Examples( | |
examples=[ | |
["An astronaut riding a sea turtle, hyper-realistic, award-winning, advertisement, 4K HD"], | |
["Dragon tiger, Victorian art style"], | |
["A dark tyranids mecha gundam style"], | |
["A dog is jumping to catch the flower"], | |
["A hobbit riding a train in a police station, digital art, highly detailed"], | |
["Female half-elf druid"], | |
["Donald Trump mixed up with Superman's suit, animation avatar style, extremely realistic"], | |
["The orc wearing a gray hat is reading a book"], | |
["The policewoman with a gas mask"], | |
["Medusa wearing a sunglass and shopping, with a snake around her neck"], | |
["Goku playing chess"], | |
["A goblin driving a snowmobile in a cave, movie poster, highly detailed"], | |
["Godzilla roaring to the sky"], | |
["Grandma is kissing a baby, detailed, (Renaissance style)"], | |
["20 year old Serbian with brown curly mullet in Naruto art form"], | |
["A beautiful elf in the world of Warcraft is drinking beer, the eye is attractive, highly detailed, animation style"], | |
["A bearded professional bald poker player holding two cards by Ron English"], | |
["A bear in red wearing a tuxedo Mr Gatsby style, hyper realistic"], | |
["A beautiful elf woman in a cyberpunk setting"], | |
["A portrait of Michael Jordan sunlight"], | |
["A beautiful steampunk warrior woman 30 years old, ultra-realistic, 8K HDR full"], | |
["A black Dragonborn Bard that plays an ocarina in a fantasy setting"], | |
["A black leather jacket glowing and floating in a medieval castle, renaissance painting, vintage"], | |
["A blue and white photograph of an old church"], | |
["A cute baby batman, 4K realistic, full body"], | |
["A dark lord grabs the soul of a man"], | |
["A demonic shaman warrior in the style of artist Frank Frazetta"], | |
["A draugr from Skyrim, dramatic lighting, cinematic photograph"], | |
["A fantasy version of Captain America wearing white armor"], | |
["Spider-Man mixed up with Hulk"], | |
["A goblin robot with metal skin, screen on its chest, drinking oil, vintage portrait, award-winning"], | |
["A hamster wearing a top hat and suit, imagining kicking a football, award-winning, realistic painting"], | |
["A happy duck with a collar, smiling, closeup, professional"], | |
["A happy moment as a bearded man with a bald head finds a key"], | |
["A He-Man figure as a futuristic warrior, retro 1980s, realistic product photography"], | |
["A hobbit with silver hair planting raspberries in a cafeteria, graffiti art, highly detailed"], | |
["A hobbit with red hair holding a compass in a plain portrait, award-winning"], | |
["A hobbit with silver hair planting trees in a rock concert sketch, award-winning"], | |
["A last alliance of men and elves marched against the armies of Mordor"], | |
["A beagle wearing a red waistcoat, blue tie, and bell, laughing out loud on the beach"], | |
["A brave man teaching an elephant to snowboard, cyberpunk style, close portrait"], | |
["A monkey with a white suit and ski goggles, laughing, surfing in the sea, hyper-realistic, award-winning, animation style"], | |
["A Nephilim attacks the city of Uruk"], | |
["Arnold Schwarzenegger shirt suit, shirtless muscle"], | |
["Armor dress style of outsider zone fantasy feathers"], | |
["Cerebro from X-Men as an unexplored wilderness, comforting colors, peaceful, inspiring"], | |
["Cerebro from X-Men as an unexplored wilderness rather than a machine"], | |
["Chad chin Sasuke, full body, detailed surrealism"], | |
["Naruto doing homework and sad, animation style, detailed portrait, painting"], | |
["Cleopatra wearing VR glasses and earphones, typing on a laptop, white dress, animation style, cyberpunk"], | |
["Dante from Devil May Cry dressed in tactical gear, realistic, full body pose"], | |
["Dark fantasy art, a well-dressed human wizard, background ancient library"], | |
["Death god wearing a cloak, playing video games"], | |
["DayZ video game bear attack scene"], | |
["Dinosaur in New York by Jean Dubuffet"], | |
["Dungeons and Dragons by Arthur Sarnoff and Dean Ellis"], | |
["Dungeons and Dragons College of Whispers Bard, Changeling male holding a Venetian mask"], | |
["Dungeons and Dragons Bugbear Merchant, fat, many ring piercings, creepy smile"], | |
["Dungeons and Dragons autumn elf"], | |
["Dungeons and Dragons anime death knight riding a fire horse"], | |
["Dungeons and Dragons dwarf paladin"], | |
["Dungeons and Dragons effeminate dwarf in pink clothes, running away with a terrified face"], | |
["Dragon Ball Super's Goku, photorealistic, ultra-detailed, 8K"], | |
["Early 1960s astronaut in the style of Alexander"], | |
["Eevee evolution in ghost type, Eevee as ghost type Pokemon, flying ghost"], | |
["Elf knight order in fantasy setting"], | |
["A portrait of Steve Jobs with thick eyeglasses, smiling, bearded, black T-shirt, watercolor painting"], | |
["Beautiful Elsa princess eating ice cream in a snowy wonderland, fantasy style, hyper-realistic"], | |
["Emperor brother, dark fantasy battle advisor, glamour portrait"], | |
["Enel from One Piece"], | |
["Eren's mom as a titan"], | |
["Eye Spy Eagle"], | |
["Fantasy female sailor, teal magic aura, large ship, lightly armored, no hat, spectral tentacles"], | |
["Fantasy guardian who feels empathy, generosity, tolerance, resilience, contemplative mood, full body view"], | |
["Fantasy RPG style portrait of a young elven mage"], | |
["Fantasy style weird human character, full body"], | |
["Fantasy strong half-orc sorcerer wearing brown robes, wild magic, background medieval farmland"], | |
["Fat Australian woman, penguin island, grotesque"], | |
["Fearsome ancient Spartan warrior from Greece, red cloak, Spartan helmet, night"], | |
["Female half-elf rogue, red hair, slightly pointed ears, canyon landscape"], | |
["Female character customization screen, RPG video game UI, illustration details"], | |
["Female beauty by the standards of 5th century Europe"], | |
["Female pirate lord, cutlass, devil eyes, 4K animated"], | |
["Female robot trooper, augmented exoskeleton, urban environment, tan leather and magnesium, fashion photography, medium shot, Nikon FX"], | |
["Fennec fox coloring book, black and white"], | |
["Film still, Gisele Bundchen as a Targaryen queen, epic shot, mid-body"], | |
["Formula 1 view from the side, off-road wheels, rugged feel, white background"], | |
["Ghost on skateboard, cartoon style"], | |
["Guardians of the Galaxy fighting in a movie theater"], | |
["Hand with glove, vector"], | |
["Happy baby mouse flying first class, black and white, for coloring"], | |
["Happy baby monkey astronaut, black and white, for coloring"], | |
["Happy pregnant black woman"], | |
["Happy steampunk style psychedelic monkey with dreadlocks, logo"], | |
["Harry Potter riding a giraffe in a secret hideaway, realistic photograph, hyper-realistic"], | |
["Henry Cavill as a gladiator, Maximus fighting"], | |
["Heraldry, a shield with a silver tower on it"], | |
["Hero figure, ancient battlefield, 3D HD"], | |
["High heels made out of fabric in a zen garden, renaissance painting, closeup"], | |
["Hyper-realistic Japanese dragon, full 8K"], | |
["Jared Leto's Joker in the style of The Batman Animated Series, episode screencapture"], | |
["Jesus raptor, hyper-realistic, white background, smiling"], | |
["The batman is eating noodles"], | |
["The giant monster roaring to the sky, highly detailed, disaster"], | |
], | |
inputs=[prompt], | |
label="Example Prompts" | |
) | |
with gr.Column(scale=1): | |
output_model = gr.Model3D( | |
label="Generated 3D Mesh", | |
camera_position=(90, 90, 3), | |
clear_color=(0.5, 0.5, 0.5, 1), | |
) | |
output_file = gr.File(label="Download OBJ file") | |
generate_btn.click( | |
fn=generate_3d_mesh, | |
inputs=[prompt], | |
outputs=[output_model, output_file] | |
) | |
gr.Markdown( | |
""" | |
## About | |
This demo uses TriplaneTurbo, which adapts Stable Diffusion for instant text-to-mesh generation. | |
The model can generate high-quality 3D meshes from text descriptions without requiring 3D training data. | |
### Limitations | |
- Generation is deterministic with a fixed seed | |
- Complex prompts may produce unpredictable results | |
- Generated meshes may require clean-up for professional use | |
""" | |
) | |
if __name__ == "__main__": | |
demo.launch() | |