| 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 |
| |
| class Model: |
| 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 run_image( |
| self, image: PIL.Image.Image, seed: int = 0, guidance_scale: float = 3.0, num_steps: int = 64 |
| ) -> str: |
| generator = torch.Generator(device=self.device).manual_seed(seed) |
| images = self.pipe_img( |
| image, |
| generator=generator, |
| guidance_scale=guidance_scale, |
| num_inference_steps=num_steps, |
| 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) |
|
|
