tex3 / app.py
hansyan's picture
Update app.py
14b375b verified
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. Here, we plug the PeRFlow-delta-weights of SD-v1.5 into the Disney-Pixar-Cartoon dreambooth.
"""
)
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 husky dog")
seed = gr.Textbox(label="Random Seed", value=42)
gr.Markdown(
"""
Images should be generated within 1 second normally, sometimes, it could a bit slow due to warm-up of the program. Here are some examples provided:
- a policeman
- a robot, close-up
- a red car, side view
- a blue mug
- a burger
- a tea pot
- a wooden chair
- a 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()