#!/usr/bin/env python import os import shutil import tempfile import gradio as gr from PIL import Image import numpy as np from settings import ( DEFAULT_IMAGE_RESOLUTION, DEFAULT_NUM_IMAGES, MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES, MAX_SEED, ) from utils import randomize_seed_fn # ---- helper to build a quick textured copy of the mesh --------------- def apply_texture(src_mesh:str, texture:str, tag:str)->str: """ Writes a copy of `src_mesh` and tiny .mtl that points to `texture`. Returns the new OBJ/GLB path for viewing. """ tmp_dir = tempfile.mkdtemp() mesh_copy = os.path.join(tmp_dir, f"{tag}.obj") mtl_name = f"{tag}.mtl" # copy geometry shutil.copy(src_mesh, mesh_copy) # write minimal MTL with open(os.path.join(tmp_dir, mtl_name), "w") as f: f.write(f"newmtl material_0\nmap_Kd {os.path.basename(texture)}\n") # ensure texture lives next to OBJ shutil.copy(texture, os.path.join(tmp_dir, os.path.basename(texture))) # patch OBJ to reference our new MTL with open(mesh_copy, "r+") as f: lines = f.readlines() if not lines[0].startswith("mtllib"): lines.insert(0, f"mtllib {mtl_name}\n") f.seek(0); f.writelines(lines) return mesh_copy def image_to_temp_path(img_like, tag): """ Convert various image-like objects (str, PIL.Image, list, tuple) to temp PNG path. Returns the path to the saved image file. """ # Handle tuple or list input if isinstance(img_like, (list, tuple)): if len(img_like) == 0: raise ValueError("Empty image list/tuple.") img_like = img_like[0] # If it's already a file path if isinstance(img_like, str): return img_like # If it's a PIL Image if isinstance(img_like, Image.Image): temp_path = os.path.join(tempfile.mkdtemp(), f"{tag}.png") img_like.save(temp_path) return temp_path # if it's numpy array if isinstance(img_like, np.ndarray): temp_path = os.path.join(tempfile.mkdtemp(), f"{tag}.png") img_like = Image.fromarray(img_like) img_like.save(temp_path) return temp_path raise ValueError(f"Expected PIL.Image, str, list, or tuple — got {type(img_like)}") def show_mesh(which, mesh, inp, coarse, fine): """Switch the displayed texture based on dropdown change.""" print() tex_map = { "Input": image_to_temp_path(inp, "input"), "Coarse": coarse[0] if isinstance(coarse, tuple) else coarse, "Fine": fine[0] if isinstance(fine, tuple) else fine, } texture_path = tex_map[which] return apply_texture(mesh, texture_path, which.lower()) # ---------------------------------------------------------------------- def create_demo(process): with gr.Blocks() as demo: with gr.Row(): with gr.Column(): image = gr.Image() prompt = gr.Textbox(label="Prompt", submit_btn=True) with gr.Accordion("Advanced options", open=False): num_samples = gr.Slider( label="Number of images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1 ) image_resolution = gr.Slider( label="Image resolution", minimum=256, maximum=MAX_IMAGE_RESOLUTION, value=DEFAULT_IMAGE_RESOLUTION, step=256, ) num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=10, step=1) guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1) seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed") n_prompt = gr.Textbox( label="Negative prompt", value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality", ) with gr.Column(): result_coarse = gr.Gallery(label="Output Coarse", show_label=True, columns=2, object_fit="scale-down") result_fine = gr.Gallery(label="Output Fine", show_label=True, columns=2, object_fit="scale-down") # mesh_viewer = gr.Model3D(label="Textured Mesh", clear_color=[0, 0, 0, 0], value="examples/monkey/mesh.obj") # radio buttons let the user toggle which texture to view # texture_choice = gr.Radio(["Input", "Coarse", "Fine"], label="Preview texture", value="Input") # mesh_path_state = gr.State("examples/bunny/mesh.obj") inputs = [ image, prompt, a_prompt, n_prompt, num_samples, image_resolution, num_steps, guidance_scale, seed, ] # first call → run diffusion / texture network prompt.submit( fn=randomize_seed_fn, inputs=[seed, randomize_seed], outputs=seed, queue=False, api_name=False, ).then( fn=process, inputs=inputs, outputs=[result_coarse, result_fine], api_name="canny", concurrency_id="main", ) # .then( # fn=show_mesh, # inputs=[texture_choice, mesh_path_state, image, result_coarse, result_fine], # outputs=mesh_viewer, # queue=False, # api_name=False, # ) gr.Examples( fn=process, inputs=inputs, outputs=[result_coarse, result_fine], examples=[ [ "examples/bunny/uv_normal.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/bunny/uv_normal/fused.png "feather", a_prompt.value, n_prompt.value, num_samples.value, image_resolution.value, num_steps.value, guidance_scale.value, seed.value, ], [ "examples/monkey/uv_normal.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/monkey/uv_normal/fused.png "wood", a_prompt.value, n_prompt.value, num_samples.value, image_resolution.value, num_steps.value, guidance_scale.value, seed.value, ], ], ) return demo if __name__ == "__main__": from model import Model model = Model(task_name="Texnet") demo = create_demo(model.process_texnet) demo.queue().launch()