ttttdiva/Shap-E

1import tempfile2 3import numpy as np4import torch5import trimesh6from shap_e.diffusion.gaussian_diffusion import diffusion_from_config7from shap_e.diffusion.sample import sample_latents8from shap_e.models.download import load_config, load_model9from shap_e.models.nn.camera import (DifferentiableCameraBatch,10                                     DifferentiableProjectiveCamera)11from shap_e.models.transmitter.base import Transmitter, VectorDecoder12from shap_e.rendering.torch_mesh import TorchMesh13from shap_e.util.collections import AttrDict14from shap_e.util.image_util import load_image15 16 17# Copied from https://github.com/openai/shap-e/blob/d99cedaea18e0989e340163dbaeb4b109fa9e8ec/shap_e/util/notebooks.py#L15-L4218def create_pan_cameras(size: int,19                       device: torch.device) -> DifferentiableCameraBatch:20    origins = []21    xs = []22    ys = []23    zs = []24    for theta in np.linspace(0, 2 * np.pi, num=20):25        z = np.array([np.sin(theta), np.cos(theta), -0.5])26        z /= np.sqrt(np.sum(z**2))27        origin = -z * 428        x = np.array([np.cos(theta), -np.sin(theta), 0.0])29        y = np.cross(z, x)30        origins.append(origin)31        xs.append(x)32        ys.append(y)33        zs.append(z)34    return DifferentiableCameraBatch(35        shape=(1, len(xs)),36        flat_camera=DifferentiableProjectiveCamera(37            origin=torch.from_numpy(np.stack(origins,38                                             axis=0)).float().to(device),39            x=torch.from_numpy(np.stack(xs, axis=0)).float().to(device),40            y=torch.from_numpy(np.stack(ys, axis=0)).float().to(device),41            z=torch.from_numpy(np.stack(zs, axis=0)).float().to(device),42            width=size,43            height=size,44            x_fov=0.7,45            y_fov=0.7,46        ),47    )48 49 50# Copied from https://github.com/openai/shap-e/blob/8625e7c15526d8510a2292f92165979268d0e945/shap_e/util/notebooks.py#LL64C1-L76C3351@torch.no_grad()52def decode_latent_mesh(53    xm: Transmitter | VectorDecoder,54    latent: torch.Tensor,55) -> TorchMesh:56    decoded = xm.renderer.render_views(57        AttrDict(cameras=create_pan_cameras(58            2, latent.device)),  # lowest resolution possible59        params=(xm.encoder if isinstance(xm, Transmitter) else60                xm).bottleneck_to_params(latent[None]),61        options=AttrDict(rendering_mode='stf', render_with_direction=False),62    )63    return decoded.raw_meshes[0]64 65 66class Model:67    def __init__(self):68        self.device = torch.device(69            'cuda' if torch.cuda.is_available() else 'cpu')70        self.xm = load_model('transmitter', device=self.device)71        self.diffusion = diffusion_from_config(load_config('diffusion'))72        self.model_text = None73        self.model_image = None74 75    def load_model(self, model_name: str) -> None:76        assert model_name in ['text300M', 'image300M']77        if model_name == 'text300M' and self.model_text is None:78            self.model_text = load_model(model_name, device=self.device)79        elif model_name == 'image300M' and self.model_image is None:80            self.model_image = load_model(model_name, device=self.device)81 82    def to_glb(self, latent: torch.Tensor) -> str:83        ply_path = tempfile.NamedTemporaryFile(suffix='.ply',84                                               delete=False,85                                               mode='w+b')86        decode_latent_mesh(self.xm, latent).tri_mesh().write_ply(ply_path)87 88        mesh = trimesh.load(ply_path.name)89        rot = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])90        mesh = mesh.apply_transform(rot)91        rot = trimesh.transformations.rotation_matrix(np.pi, [0, 1, 0])92        mesh = mesh.apply_transform(rot)93 94        mesh_path = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)95        mesh.export(mesh_path.name, file_type='glb')96 97        return mesh_path.name98 99    def run_text(self,100                 prompt: str,101                 seed: int = 0,102                 guidance_scale: float = 15.0,103                 num_steps: int = 64) -> str:104        self.load_model('text300M')105        torch.manual_seed(seed)106 107        latents = sample_latents(108            batch_size=1,109            model=self.model_text,110            diffusion=self.diffusion,111            guidance_scale=guidance_scale,112            model_kwargs=dict(texts=[prompt]),113            progress=True,114            clip_denoised=True,115            use_fp16=True,116            use_karras=True,117            karras_steps=num_steps,118            sigma_min=1e-3,119            sigma_max=160,120            s_churn=0,121        )122        return self.to_glb(latents[0])123 124    def run_image(self,125                  image_path: str,126                  seed: int = 0,127                  guidance_scale: float = 3.0,128                  num_steps: int = 64) -> str:129        self.load_model('image300M')130        torch.manual_seed(seed)131 132        image = load_image(image_path)133        latents = sample_latents(134            batch_size=1,135            model=self.model_image,136            diffusion=self.diffusion,137            guidance_scale=guidance_scale,138            model_kwargs=dict(images=[image]),139            progress=True,140            clip_denoised=True,141            use_fp16=True,142            use_karras=True,143            karras_steps=num_steps,144            sigma_min=1e-3,145            sigma_max=160,146            s_churn=0,147        )148        return self.to_glb(latents[0])149