File size: 3,407 Bytes
1025344 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import tempfile
import numpy as np
import PIL.Image
import torch
import trimesh
from diffusers import ShapEImg2ImgPipeline, ShapEPipeline
from diffusers.utils import export_to_ply, export_to_gif
from typing import Dict, List, Any
class EndpointHandler():
def __init__(self):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.pipe = ShapEPipeline.from_pretrained("openai/shap-e", torch_dtype=torch.float16)
self.pipe.to(self.device)
self.pipe_img = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", torch_dtype=torch.float16)
self.pipe_img.to(self.device)
def to_glb(self, ply_path: str) -> str:
mesh = trimesh.load(ply_path)
rot = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])
mesh = mesh.apply_transform(rot)
rot = trimesh.transformations.rotation_matrix(np.pi, [0, 1, 0])
mesh = mesh.apply_transform(rot)
mesh_path = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
mesh.export(mesh_path.name, file_type="glb")
return mesh_path.name
def run_text(self, prompt: str, seed: int = 0, guidance_scale: float = 15.0, num_steps: int = 64, output_type: str = "pil") -> str:
generator = torch.Generator(device=self.device).manual_seed(seed)
if output_type=="pil":
images = self.pipe(
prompt,
num_images_per_prompt=4,
generator=generator,
guidance_scale=guidance_scale,
num_inference_steps=num_steps,
frame_size=64,
).images
gif_path = tempfile.NamedTemporaryFile(suffix=".gif", delete=False, mode="w+b")
gif_path2 = tempfile.NamedTemporaryFile(suffix=".gif", delete=False, mode="w+b")
gif_path3 = tempfile.NamedTemporaryFile(suffix=".gif", delete=False, mode="w+b")
gif_path4 = tempfile.NamedTemporaryFile(suffix=".gif", delete=False, mode="w+b")
export_to_gif(images[0], gif_path.name)
export_to_gif(images[1], gif_path2.name)
export_to_gif(images[2], gif_path3.name)
export_to_gif(images[3], gif_path4.name)
return gif_path.name, gif_path2.name, gif_path3.name, gif_path4.name
else:
images = self.pipe(
prompt,
num_images_per_prompt=1,
generator=generator,
guidance_scale=guidance_scale,
num_inference_steps=num_steps,
frame_size=64,
output_type=output_type
).images
ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False, mode="w+b")
export_to_ply(images[0], ply_path.name)
return self.to_glb(ply_path.name)
def __call__(self, prompt: str, seed: int = 0, guidance_scale: float = 15.0, num_steps: int = 64) -> str:
generator = torch.Generator(device=self.device).manual_seed(seed)
images = self.pipe(
prompt,
num_images_per_prompt=1,
generator=generator,
guidance_scale=guidance_scale,
num_inference_steps=num_steps,
frame_size=64,
output_type="mesh"
).images
ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False, mode="w+b")
export_to_ply(images[0], ply_path.name)
return self.to_glb(ply_path.name) |