import spaces import os, logging, time, argparse, random, tempfile, rembg, shlex, subprocess import gradio as gr import numpy as np import torch from PIL import Image from functools import partial subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl')) from tsr.system import TSR from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation from src.scheduler_perflow import PeRFlowScheduler from diffusers import StableDiffusionPipeline, UNet2DConditionModel def fill_background(img): img = np.array(img).astype(np.float32) / 255.0 img = img[:, :, :3] * img[:, :, 3:4] + (1 - img[:, :, 3:4]) * 0.5 img = Image.fromarray((img * 255.0).astype(np.uint8)) return img def merge_delta_weights_into_unet(pipe, delta_weights, org_alpha = 1.0): unet_weights = pipe.unet.state_dict() for key in delta_weights.keys(): dtype = unet_weights[key].dtype try: unet_weights[key] = org_alpha * unet_weights[key].to(dtype=delta_weights[key].dtype) + delta_weights[key].to(device=unet_weights[key].device) except: unet_weights[key] = unet_weights[key].to(dtype=delta_weights[key].dtype) unet_weights[key] = unet_weights[key].to(dtype) pipe.unet.load_state_dict(unet_weights, strict=True) return pipe def setup_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True if torch.cuda.is_available(): device = "cuda:0" else: device = "cpu" ### TripoSR model = TSR.from_pretrained( "stabilityai/TripoSR", config_name="config.yaml", weight_name="model.ckpt", ) # adjust the chunk size to balance between speed and memory usage model.renderer.set_chunk_size(8192) model.to(device) ### PeRFlow-T2I # pipe_t2i = StableDiffusionPipeline.from_pretrained("Lykon/dreamshaper-8", torch_dtype=torch.float16, safety_checker=None) # pipe_t2i = StableDiffusionPipeline.from_pretrained("stablediffusionapi/disney-pixar-cartoon", torch_dtype=torch.float16, safety_checker=None) # delta_weights = UNet2DConditionModel.from_pretrained("hansyan/piecewise-rectified-flow-delta-weights", torch_dtype=torch.float16, variant="v0-1",).state_dict() # pipe_t2i = merge_delta_weights_into_unet(pipe_t2i, delta_weights) pipe_t2i = StableDiffusionPipeline.from_pretrained("hansyan/perflow-sd15-disney", torch_dtype=torch.float16, safety_checker=None) pipe_t2i.scheduler = PeRFlowScheduler.from_config(pipe_t2i.scheduler.config, prediction_type="epsilon", num_time_windows=4) pipe_t2i.to('cuda:0', torch.float16) ### gradio rembg_session = rembg.new_session() @spaces.GPU def generate(text, seed): def fill_background(image): image = np.array(image).astype(np.float32) / 255.0 image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5 image = Image.fromarray((image * 255.0).astype(np.uint8)) return image setup_seed(int(seed)) prompt_prefix = "high quality, highly detailed, (best quality, masterpiece), " neg_prompt = "EasyNegative, drawn by bad-artist, sketch by bad-artist-anime, (bad_prompt:0.8), (artist name, signature, watermark:1.4), (ugly:1.2), (worst quality, poor details:1.4), bad-hands-5, badhandv4, blurry" text = prompt_prefix + text samples = pipe_t2i( prompt = [text], negative_prompt = [neg_prompt], height = 512, width = 512, # num_inference_steps = 6, # guidance_scale = 7.5, num_inference_steps = 8, guidance_scale = 7.5, output_type = 'pt', ).images samples = samples.squeeze(0).permute(1, 2, 0).cpu().numpy()*255. samples = samples.astype(np.uint8) samples = Image.fromarray(samples[:, :, :3]) return samples @spaces.GPU def render(image, mc_resolution=256, formats=["obj"]): image = Image.fromarray(image) image = image.resize((768, 768)) image = remove_background(image, rembg_session) image = resize_foreground(image, 0.85) image = fill_background(image) scene_codes = model(image, device=device) mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0] mesh = to_gradio_3d_orientation(mesh) rv = [] for format in formats: mesh_path = tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False) mesh.export(mesh_path.name) rv.append(mesh_path.name) return rv[0] # layout css = """ h1 { text-align: center; display:block; } h2 { text-align: center; display:block; } h3 { text-align: center; display:block; } """ with gr.Blocks(title="TripoSR", css=css) as interface: gr.Markdown( """ # Instant Text-to-3D Mesh Demo ### [PeRFlow](https://github.com/magic-research/piecewise-rectified-flow)-T2I + [TripoSR](https://github.com/VAST-AI-Research/TripoSR) Two-stage synthesis: 1) generating images by PeRFlow-T2I; 2) rendering 3D assests. """ ) with gr.Column(): with gr.Row(): output_image = gr.Image(label='Generated Image', height=384,) output_model_obj = gr.Model3D( label="Output 3D Model (OBJ Format)", interactive=False, height=384, ) with gr.Row(): textbox = gr.Textbox(label="Input Prompt", value="a colorful bird") seed = gr.Textbox(label="Random Seed", value=42) gr.Markdown( """ Examples: - a policeman - a robot, close-up - a red car, side view - a blue mug - a burger - a tea pot - a wooden chair - an amazing unicorn """ ) # activate textbox.submit( fn=generate, inputs=[textbox, seed], outputs=[output_image], ).success( fn=render, inputs=[output_image], outputs=[output_model_obj], ) seed.submit( fn=generate, inputs=[textbox, seed], outputs=[output_image], ).success( fn=render, inputs=[output_image], outputs=[output_model_obj], ) if __name__ == '__main__': interface.queue(max_size=10) interface.launch()