import os import sys import spaces import base64 import tempfile import hashlib import re import uuid import shutil from omegaconf import OmegaConf from typing import Optional, Union, Tuple import gradio as gr GUIDEFLOW_YELLOW = "#ccad57" GUIDEFLOW_BLUE = "#2459c2" GUIDEFLOW_GREEN = "#8edf9f" os.environ["CUMM_DISABLE_JIT"] = "1" os.environ["SPCONV_DISABLE_JIT"] = "1" os.environ["TOKENIZERS_PARALLELISM"] = "false" # Add project root to Python path project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if project_root not in sys.path: sys.path.insert(0, project_root) # --- START XVFB GLOBALLY --- # Check if we are in a headless environment and DISPLAY is not set if os.environ.get("DISPLAY") is None: print("[INFO] Starting Xvfb for headless rendering...") from pyvirtualdisplay import Display # Start Xvfb. visible=0 means headless. # size=(1920, 1080) matches your previous xvfb-run settings. display = Display(visible=0, size=(1920, 1080)) display.start() # Ensure DISPLAY env var is set for subprocesses if os.environ.get("DISPLAY") is None: # PyVirtualDisplay usually sets this, but fallback if needed os.environ["DISPLAY"] = f":{display.display}" print(f"[INFO] Xvfb started on {os.environ['DISPLAY']}") # --- LOGO SETUP (BASE64) --- def image_to_base64(image_path): """Encodes an image to a base64 string for direct HTML embedding.""" if not os.path.exists(image_path): return "" with open(image_path, "rb") as img_file: encoded_string = base64.b64encode(img_file.read()).decode('utf-8') return f"data:image/png;base64,{encoded_string}" logo_rel_path = os.path.join("demos", "assets", "logo.png") logo_abs_path = os.path.join(project_root, logo_rel_path) logo_src = image_to_base64(logo_abs_path) BLENDER_LINK = 'https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz' BLENDER_INSTALLATION_PATH = '/tmp' BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64/blender' def _install_blender(): if not os.path.exists(BLENDER_PATH): os.system('sudo apt-get update') os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6') os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}') os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}') def _download_objaverse_ckpt(): if not os.path.exists(os.path.join(project_root, 'models', 'model_objaverse.ckpt')): os.makedirs(os.path.join(project_root, 'models'), exist_ok=True) os.system(f'wget https://huggingface.co/mikaelaangel/partfield-ckpt/resolve/main/model_objaverse.ckpt -O {os.path.join(project_root, 'models', 'model_objaverse.ckpt')}') _install_blender() _download_objaverse_ckpt() # Attempt import, handle failure gracefully for the demo shell try: from demos.pipeline_fn import GuideFlow3dPipeline except ImportError: GuideFlow3dPipeline = None pipe = None cfg = None # Initialize Pipeline try: cfg_path = os.path.join(project_root, 'config', 'default.yaml') if os.path.exists(cfg_path): cfg = OmegaConf.load(cfg_path) if GuideFlow3dPipeline: pipe = GuideFlow3dPipeline().from_pretrained(cfg) except Exception as e: print(f"Error initializing pipeline: {e}") pass output_dir = os.path.join(project_root, "all_outputs") os.makedirs(output_dir, exist_ok=True) # --- MAPPING HELPERS --- # Dictionary mapping static thumbnail images to actual GLB files THUMB_TO_GLB = { # Structure Mesh Examples "example_data/thumbs/structure/bench_chair.png": "example_data/structure_mesh/bench_chair.glb", "example_data/thumbs/structure/cabinet.png": "example_data/structure_mesh/cabinet.glb", "example_data/thumbs/structure/chair.png": "example_data/structure_mesh/chair.glb", "example_data/thumbs/structure/giraffe.png": "example_data/structure_mesh/giraffe.glb", "example_data/thumbs/structure/motorcycle.png": "example_data/structure_mesh/motorcycle.glb", "example_data/thumbs/structure/plane.png": "example_data/structure_mesh/plane.glb", # Reference Appearance Mesh Examples "example_data/thumbs/appearance/B01DA8LC0A.jpg": "example_data/appearance_mesh/B01DA8LC0A.glb", "example_data/thumbs/appearance/B01DJH73Y6.png": "example_data/appearance_mesh/B01DJH73Y6.glb", "example_data/thumbs/appearance/B0728KSP33.jpg": "example_data/appearance_mesh/B0728KSP33.glb", "example_data/thumbs/appearance/B07B4YXNR8.jpg": "example_data/appearance_mesh/B07B4YXNR8.glb", "example_data/thumbs/appearance/B07QC84LP1.png": "example_data/appearance_mesh/B07QC84LP1.glb", "example_data/thumbs/appearance/B07QFRSC8M.png": "example_data/appearance_mesh/B07QFRSC8M_zup.glb", "example_data/thumbs/appearance/B082QC7YKR.png": "example_data/appearance_mesh/B082QC7YKR_zup.glb" } # Create a lookup based on basename to be robust against Gradio temp paths THUMB_BASENAME_TO_GLB = {os.path.basename(k): v for k, v in THUMB_TO_GLB.items()} # Create reverse lookup for strict example detection GLB_ABS_PATH_TO_NAME = {} for k, v in THUMB_TO_GLB.items(): abs_p = os.path.abspath(os.path.join(project_root, v)) name_no_ext = os.path.splitext(os.path.basename(v))[0] GLB_ABS_PATH_TO_NAME[abs_p] = name_no_ext def load_mesh_from_thumb(thumb_path: str) -> Optional[str]: """Callback to return the GLB path associated with a thumbnail.""" if not thumb_path: return None basename = os.path.basename(thumb_path) return THUMB_BASENAME_TO_GLB.get(basename, None) def _ensure_glb_path(result: Union[str, bytes, os.PathLike]) -> str: """Normalize various return types from fn() to a .glb file path.""" if isinstance(result, (str, os.PathLike)): path = os.fspath(result) if not os.path.exists(path): raise gr.Error("Returned mesh path does not exist.") return path if isinstance(result, (bytes, bytearray)): tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".glb") tmp.write(result) tmp.flush() tmp.close() return tmp.name def file_sha256(path: str, chunk_size: int = 1 << 20) -> str: h = hashlib.sha256() if not path or not os.path.exists(path): return "nocontent" with open(path, "rb") as f: for chunk in iter(lambda: f.read(chunk_size), b""): h.update(chunk) return h.hexdigest() def get_cache_folder(struct_mesh_path: str) -> str: """Determines the output folder name based on the structure mesh.""" struct_abs = os.path.abspath(struct_mesh_path) # 1. Check if it is a known example # We check both absolute path and if the file hash matches a known example (to be safe) # But simplest is just path check for now as examples are static files. if struct_abs in GLB_ABS_PATH_TO_NAME: return GLB_ABS_PATH_TO_NAME[struct_abs] # Also check if basename matches an example (for when gradio moves files around but keeps names?) # Actually, safely relying on content hash is better for "New" files. current_hash = file_sha256(struct_mesh_path) # 2. Scan existing temp_* folders for matching hash # We look for a struct_mesh.hash file inside. if os.path.exists(output_dir): for item in os.listdir(output_dir): if item.startswith("temp_"): folder_path = os.path.join(output_dir, item) if os.path.isdir(folder_path): hash_file = os.path.join(folder_path, "struct_mesh.hash") if os.path.exists(hash_file): try: with open(hash_file, "r") as f: stored_hash = f.read().strip() if stored_hash == current_hash: return item except: continue # 3. If not found, create new temp_{id} new_id = uuid.uuid4().hex[:8] return f"temp_{new_id}" # @spaces.GPU(duration=360) def on_run( guidance_mode_state: str, app_struct_mesh: Optional[str], app_ref_mesh: Optional[str], app_ref_image: Optional[str], sim_struct_mesh: Optional[str], sim_ref_text: Optional[str], sim_ref_image: Optional[str], target_up_label: str, reference_up_label: str, cfg_strength: float, num_steps: int, learning_rate: float, ) -> Tuple[str, Optional[str]]: current_mode = guidance_mode_state.lower() if current_mode == "appearance": target_mesh_path = app_struct_mesh reference_mesh_path = app_ref_mesh reference_image_path = app_ref_image reference_text = None else: target_mesh_path = sim_struct_mesh reference_text = sim_ref_text reference_image_path = sim_ref_image reference_mesh_path = None if not target_mesh_path: raise gr.Error(f"Target Structure mesh is required for {current_mode} mode.") if pipe is None: raise gr.Error("Pipeline not initialized. Check logs.") # --- Determine Output Directory --- folder_name = get_cache_folder(target_mesh_path) run_output_dir = os.path.join(output_dir, folder_name) os.makedirs(run_output_dir, exist_ok=True) print(f"[INFO] Using output directory: {run_output_dir}") args = { "structure_mesh": target_mesh_path, "output_dir": run_output_dir, "convert_target_yup_to_zup": target_up_label == "Z-up", "convert_appearance_yup_to_zup": reference_up_label == "Z-up", "appearance_mesh": reference_mesh_path, "appearance_image": reference_image_path, "appearance_text": (reference_text or "").strip(), } fn = None if current_mode == "appearance": if not reference_mesh_path: raise gr.Error("Appearance mode requires a reference mesh.") fn = pipe.run_appearance args.pop("appearance_text", None) else: # similarity if not reference_text and reference_image_path: args["appearance_image"] = reference_image_path args.pop("appearance_text", None) args["app_type"] = "image" elif reference_text and not reference_image_path: args["appearance_text"] = reference_text args.pop("appearance_image", None) args["app_type"] = "text" elif reference_text and reference_image_path: raise gr.Error("Similarity mode requires a text prompt or reference image, but not both.") else: raise gr.Error("Similarity mode requires a text prompt or reference image.") fn = pipe.run_self_similarity args.pop("appearance_mesh", None) args.pop("convert_appearance_yup_to_zup", None) if cfg: updated_cfg = cfg # OmegaConf.load(cfg) updated_cfg.cfg_strength = cfg_strength updated_cfg.steps = num_steps updated_cfg.learning_rate = learning_rate pipe.cfg = updated_cfg try: result_mesh, result_video = fn(**args) mesh_path = _ensure_glb_path(result_mesh) video_path = _ensure_glb_path(result_video) return mesh_path, video_path except Exception as e: raise gr.Error(f"Generation failed: {str(e)}") # --- UI Styling & Header --- font_reg = os.path.join(project_root, "demos", "assets", "fonts", "avenir-next", "AvenirNextCyr-Regular.ttf") font_bold = os.path.join(project_root, "demos", "assets", "fonts", "avenir-next", "AvenirNextCyr-Bold.ttf") font_heavy = os.path.join(project_root, "demos", "assets", "fonts", "avenir-next", "AvenirNextCyr-Heavy.ttf") css = f""" @font-face {{ font-family: 'Avenir Next Regular'; src: url('/file={font_reg}') format('truetype'); font-weight: normal; font-style: normal; }} @font-face {{ font-family: 'Avenir Next Bold'; src: url('/file={font_bold}') format('truetype'); font-weight: bold; font-style: normal; }} @font-face {{ font-family: 'Avenir Next Heavy'; src: url('/file={font_heavy}') format('truetype'); font-weight: normal; font-style: normal; }} body, .gradio-container {{ background-color: #ffffff !important; color: #1f2937 !important; font-family: 'Avenir Next Regular', sans-serif !important; }} .dark body, .dark .gradio-container {{ background-color: #ffffff !important; color: #1f2937 !important; font-family: 'Avenir Next Regular', 'Inter', 'Roboto', sans-serif !important; }} /* Add specific components */ .gradio-container button, .gradio-container input, .gradio-container textarea, .gradio-container label, .gradio-container span, .gradio-container p, .gradio-container h1, .gradio-container h2, .gradio-container h3, .gradio-container h4, .gradio-container h5, .gradio-container h6 {{ font-family: 'Avenir Next Regular', sans-serif !important; }} .guideflow-header {{ display: flex; flex-direction: column; align-items: center; margin-bottom: 1rem; transform: translateY(0.5rem); }} .logo-row {{ display: flex; align-items: baseline; gap: 0.2rem; }} .logo-img {{ height: 4rem; width: auto; transform: translateY(0.5rem); }} .title-uide, .title-flow, .title-3d {{ font-family: 'Avenir Next Regular', sans-serif !important; font-size: 3.5rem; font-weight: normal; line-height: 1.2; }} .title-uide {{ background: linear-gradient(90deg, {GUIDEFLOW_GREEN}, {GUIDEFLOW_BLUE}); -webkit-background-clip: text; background-clip: text; color: transparent; }} .title-flow {{ color: {GUIDEFLOW_BLUE}; }} .title-3d {{ color: {GUIDEFLOW_YELLOW}; }} .subtitle {{ font-size: 1.5rem; font-family: 'Avenir Next Regular', sans-serif; color: {GUIDEFLOW_YELLOW}; margin-top: 0.5rem; text-align: center; }} .authors {{ font-size: 1rem; color: #334155; margin-top: 0.5rem; }} .affiliations {{ font-size: 0.9rem; color: #6b7280; margin-top: 0.2rem; }} .venue {{ font-size: 1.1rem; font-weight: 700; color: #111827; margin-top: 0.5rem; }} .links a {{ color: {GUIDEFLOW_BLUE}; text-decoration: none; margin: 0 0.5rem; font-weight: 500; }} .links a:hover {{ text-decoration: underline; }} .demo-credit {{ font-size: 0.9rem; color: #64748b; margin-top: 0.5rem; }} .instructions-container {{ max-width: 800px; margin: 0 auto 2rem auto; text-align: left; padding: 0 1rem; }} .input-row {{ align-items: flex-start; margin-bottom: 1rem; }} """ HEADER_HTML = f"""
GuideFlow3D Logo uideFlow3D
Optimization-Guided Rectified Flow For Appearance Transfer
Sayan Deb Sarkar1    Sinisa Stekovic2    Vincent Lepetit2    Iro Armeni1
1Stanford University    2ENPC, IP Paris
NeurIPS 2025
Paper | Project Page | GitHub
Demo made by Suvaditya Mukherjee
""" INSTRUCTIONS_MD = """

Instructions

  1. Upload a Structure Mesh (.glb): This defines the shape of your 3D object. We expect a y-up mesh, but feel free to convert it using the "Advanced Settings" below.
  2. Choose Guidance Mode: Select "Self-Similarity" (Text) or "Appearance" (Mesh/Image) using the tabs.
  3. Provide Reference: Enter a text prompt or upload a reference image/mesh.
  4. Run: Click "Generate 3D Asset" to create the result.
  5. Result: The result will be displayed in the viewer on the left, and a video will be generated on the right.
""" # Example Data EX_STRUCT_THUMBS = [ ["example_data/thumbs/structure/bench_chair.png"], ["example_data/thumbs/structure/cabinet.png"], ["example_data/thumbs/structure/chair.png"], ["example_data/thumbs/structure/giraffe.png"], ["example_data/thumbs/structure/motorcycle.png"], ["example_data/thumbs/structure/plane.png"] ] EX_MESH_THUMBS = [ ["example_data/thumbs/appearance/B01DA8LC0A.jpg"], ["example_data/thumbs/appearance/B01DJH73Y6.png"], ["example_data/thumbs/appearance/B0728KSP33.jpg"], ["example_data/thumbs/appearance/B07B4YXNR8.jpg"], ["example_data/thumbs/appearance/B07QC84LP1.png"], ["example_data/thumbs/appearance/B07QFRSC8M.png"], ["example_data/thumbs/appearance/B082QC7YKR.png"] ] EX_IMG = [ "example_data/appearance_image/B01DA8LC0A.jpg", "example_data/appearance_image/B01DJH73Y6.png", "example_data/appearance_image/B0728KSP33.jpg", "example_data/appearance_image/B07B4YXNR8.jpg", "example_data/appearance_image/B07QC84LP1.png", "example_data/appearance_image/B07QFRSC8M.jpg", "example_data/appearance_image/B082QC7YKR.png" ] EX_TEXT = ["a wooden chair", "a marble statue", "A black metal-framed bed with a curved headboard, white rectangular mattress, and two white rectangular pillows.", "Rectangular wooden cabinet with reddish-brown finish, standing on four short legs. Features two drawers (upper larger, lower smaller) and an open shelf below. Back has a power socket extension and cable, ideal for electronics."] with gr.Blocks( title="GuideFlow3D", ) as demo: gr.HTML(HEADER_HTML) gr.HTML(INSTRUCTIONS_MD) guidance_mode_state = gr.State(value="Similarity") with gr.Tabs() as guidance_tabs: # --- TAB 1: SELF-SIMILARITY (LEFT) --- with gr.TabItem("Self-Similarity", id="tab_similarity") as tab_sim: gr.Markdown("### Similarity Editing Inputs") with gr.Row(elem_classes="input-row"): with gr.Column(scale=3): sim_struct_mesh = gr.Model3D(label="Structure Mesh (.glb)", interactive=True, height=300) with gr.Column(scale=2): sim_struct_hidden = gr.Image(type="filepath", visible=False) # sim_struct_mesh_examples = gr.Examples(examples=EX_STRUCT_THUMBS, inputs=sim_struct_hidden, label="Structure Examples") sim_struct_mesh_examples = gr.Examples( examples=EX_STRUCT_THUMBS, inputs=sim_struct_hidden, outputs=sim_struct_mesh, # Target the 3D viewer directly fn=load_mesh_from_thumb, # Run the conversion function run_on_click=True, # Force execution on click label="Structure Examples" ) gr.Markdown("> **_NOTE:_** Please use either a reference image or a reference text prompt, but not both.") with gr.Row(elem_classes="input-row"): with gr.Column(scale=3): sim_ref_image = gr.Image(label="Reference Appearance Image", type="filepath", height=250) with gr.Column(scale=2): gr.Examples(examples=EX_IMG, inputs=sim_ref_image, label="Image Examples") with gr.Row(elem_classes="input-row"): with gr.Column(scale=3): sim_ref_text = gr.Textbox(label="Reference Text Prompt", placeholder="Describe the appearance...", lines=2) with gr.Column(scale=2): gr.Examples(examples=EX_TEXT, inputs=sim_ref_text, label="Prompt Examples") # --- TAB 2: APPEARANCE (RIGHT) --- with gr.TabItem("Appearance", id="tab_appearance") as tab_app: gr.Markdown("### Appearance Transfer Inputs") with gr.Row(elem_classes="input-row"): with gr.Column(scale=3): app_struct_mesh = gr.Model3D(label="Structure Mesh (.glb)", interactive=True, height=300) with gr.Column(scale=2): app_struct_hidden = gr.Image(type="filepath", visible=False) # app_struct_mesh_examples = gr.Examples(examples=EX_STRUCT_THUMBS, inputs=app_struct_hidden, label="Structure Examples") app_struct_mesh_examples = gr.Examples( examples=EX_STRUCT_THUMBS, inputs=app_struct_hidden, outputs=app_struct_mesh, # Target the 3D viewer directly fn=load_mesh_from_thumb, # Run the conversion function run_on_click=True, # Force execution on click label="Structure Examples" ) with gr.Row(elem_classes="input-row"): with gr.Column(scale=3): app_ref_image = gr.Image(label="Reference Appearance Image", type="filepath", height=250) with gr.Column(scale=2): gr.Examples(examples=EX_IMG, inputs=app_ref_image, label="Image Examples") with gr.Row(elem_classes="input-row"): with gr.Column(scale=3): app_ref_mesh = gr.Model3D(label="Reference Appearance Mesh (.glb)", interactive=True, height=300) with gr.Column(scale=2): app_ref_mesh_hidden = gr.Image(type="filepath", visible=False) # app_ref_mesh_examples = gr.Examples(examples=EX_MESH_THUMBS, inputs=app_ref_mesh_hidden, label="Mesh Examples") app_ref_mesh_examples = gr.Examples( examples=EX_MESH_THUMBS, inputs=app_ref_mesh_hidden, outputs=app_ref_mesh, # Target the 3D viewer directly fn=load_mesh_from_thumb, # Run the conversion function run_on_click=True, # Force execution on click label="Mesh Examples" ) # --- ADVANCED SETTINGS --- with gr.Accordion("Advanced Settings", open=False): with gr.Row(): target_up = gr.Radio(["Y-up", "Z-up"], value="Y-up", label="Structure Mesh Up-Axis") reference_up = gr.Radio(["Y-up", "Z-up"], value="Y-up", label="Appearance Mesh Up-Axis") with gr.Row(): cfg_strength = gr.Slider(0.1, 10.0, value=5.0, step=0.1, label="CFG Strength") num_steps = gr.Slider(50, 1000, value=300, step=50, label="Diffusion Steps") learning_rate = gr.Number(value=5e-4, label="Learning Rate") # --- RUN BUTTON --- with gr.Row(): run_btn = gr.Button("Generate 3D Asset", variant="primary", size="lg") # --- OUTPUTS --- gr.Markdown("### Results") with gr.Row(): with gr.Column(): output_model = gr.Model3D(label="Output Mesh", interactive=False, clear_color=[1.0, 1.0, 1.0, 0.0]) with gr.Column(): output_video = gr.Video(label="Output Video", autoplay=True, loop=True, interactive=False) tab_sim.select(lambda: "Similarity", outputs=guidance_mode_state) tab_app.select(lambda: "Appearance", outputs=guidance_mode_state) run_btn.click( fn=on_run, inputs=[ guidance_mode_state, app_struct_mesh, app_ref_mesh, app_ref_image, sim_struct_mesh, sim_ref_text, sim_ref_image, target_up, reference_up, cfg_strength, num_steps, learning_rate ], outputs=[output_model, output_video] ) demo.load(None, None, None, js="() => { document.body.classList.remove('dark'); }") if __name__ == "__main__": # demo.queue().launch(share=True, allowed_paths=[project_root], mcp_server=True) # Useful for Colab runs demo.queue().launch( allowed_paths=[project_root], mcp_server=True, css=css, theme=gr.themes.Default( primary_hue="sky", secondary_hue="lime" ).set( body_background_fill="white", background_fill_primary="white", block_background_fill="white", input_background_fill="#f9fafb" ) )