add handler

handler.py

@@ -0,0 +1,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)