Spaces:
Running on Zero
Running on Zero
Bobby
Claude Opus 5.5
Gated core speedups: GPU preview render, SLAT/SS SDPA+batched CFG, sync-free bake, upstream SS schedule
a6a6aa0 | import torch | |
| import numpy as np | |
| from tqdm import tqdm | |
| import utils3d | |
| from easydict import EasyDict as edict | |
| from ..renderers import GaussianRenderer, MeshRenderer | |
| from ..representations import Gaussian, MeshExtractResult | |
| # from ..modules import sparse as sp | |
| from .random_utils import sphere_hammersley_sequence | |
| def yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, rs, fovs): | |
| with torch.no_grad(): | |
| is_list = isinstance(yaws, list) | |
| if not is_list: | |
| yaws = [yaws] | |
| pitchs = [pitchs] | |
| if not isinstance(rs, list): | |
| rs = [rs] * len(yaws) | |
| if not isinstance(fovs, list): | |
| fovs = [fovs] * len(yaws) | |
| extrinsics = [] | |
| intrinsics = [] | |
| for yaw, pitch, r, fov in zip(yaws, pitchs, rs, fovs): | |
| fov = torch.deg2rad(torch.tensor(float(fov))).cuda() | |
| yaw = torch.tensor(float(yaw)).cuda() | |
| pitch = torch.tensor(float(pitch)).cuda() | |
| orig = torch.tensor([ | |
| torch.sin(yaw) * torch.cos(pitch), | |
| torch.cos(yaw) * torch.cos(pitch), | |
| torch.sin(pitch), | |
| ]).cuda() * r | |
| extr = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) | |
| intr = utils3d.torch.intrinsics_from_fov_xy(fov, fov) | |
| extrinsics.append(extr) | |
| intrinsics.append(intr) | |
| if not is_list: | |
| extrinsics = extrinsics[0] | |
| intrinsics = intrinsics[0] | |
| return extrinsics, intrinsics | |
| def get_renderer(sample, **kwargs): | |
| if isinstance(sample, Gaussian): | |
| renderer = GaussianRenderer() | |
| renderer.rendering_options.resolution = kwargs.get('resolution', 512) | |
| renderer.rendering_options.near = kwargs.get('near', 0.8) | |
| renderer.rendering_options.far = kwargs.get('far', 1.6) | |
| renderer.rendering_options.bg_color = kwargs.get('bg_color', (0, 0, 0)) | |
| renderer.rendering_options.ssaa = kwargs.get('ssaa', 1) | |
| renderer.pipe.kernel_size = kwargs.get('kernel_size', 0.1) | |
| renderer.pipe.use_mip_gaussian = True | |
| elif isinstance(sample, MeshExtractResult): | |
| renderer = MeshRenderer() | |
| renderer.rendering_options.resolution = kwargs.get('resolution', 512) | |
| renderer.rendering_options.near = kwargs.get('near', 1) | |
| renderer.rendering_options.far = kwargs.get('far', 100) | |
| renderer.rendering_options.ssaa = kwargs.get('ssaa', 1) | |
| else: | |
| raise ValueError(f'Unsupported sample type: {type(sample)}') | |
| return renderer | |
| def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=None, verbose=True, **kwargs): | |
| renderer = get_renderer(sample, **options) | |
| rets = {} | |
| mode = kwargs.get('mode') | |
| for j, (extr, intr) in tqdm(enumerate(zip(extrinsics, intrinsics)), desc='Rendering', disable=not verbose): | |
| if isinstance(sample, MeshExtractResult): | |
| has_vertex_color = getattr(sample, "vertex_attrs", None) is not None | |
| return_types = ["normal"] | |
| if mode == "color": | |
| return_types = ["color"] if has_vertex_color else ["normal"] | |
| elif mode == "normal": | |
| return_types = ["normal"] | |
| else: | |
| # Default behavior for mesh keeps both channels available. | |
| return_types = ["color", "normal"] if has_vertex_color else ["normal"] | |
| res = renderer.render(sample, extr, intr, return_types=return_types) | |
| if 'color' in res: | |
| if 'color' not in rets: | |
| rets['color'] = [] | |
| rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) | |
| if 'normal' in res: | |
| if 'normal' not in rets: | |
| rets['normal'] = [] | |
| rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) | |
| if mode == "color" and 'color' not in rets and 'normal' in rets: | |
| # Fallback to normal frames when no vertex colors exist. | |
| rets['color'] = list(rets['normal']) | |
| else: | |
| res = renderer.render(sample, extr, intr, colors_overwrite=colors_overwrite) | |
| if 'color' not in rets: rets['color'] = [] | |
| if 'depth' not in rets: rets['depth'] = [] | |
| rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) | |
| if 'percent_depth' in res: | |
| rets['depth'].append(res['percent_depth'].detach().cpu().numpy()) | |
| elif 'depth' in res: | |
| rets['depth'].append(res['depth'].detach().cpu().numpy()) | |
| else: | |
| rets['depth'].append(None) | |
| return rets | |
| def render_video(sample, resolution=512, bg_color=(0, 0, 0), num_frames=300, r=2, fov=40, **kwargs): | |
| with torch.no_grad(): | |
| yaws = torch.linspace(0, 2 * 3.1415, num_frames) | |
| pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames)) | |
| yaws = yaws.tolist() | |
| pitch = pitch.tolist() | |
| extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, r, fov) | |
| return render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs) | |
| def render_multiview(sample, resolution=512, nviews=10): | |
| r = 2 | |
| fov = 40 | |
| cams = [sphere_hammersley_sequence(i, nviews) for i in range(nviews)] | |
| yaws = [cam[0] for cam in cams] | |
| pitchs = [cam[1] for cam in cams] | |
| extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, r, fov) | |
| res = render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': (0, 0, 0)}) | |
| if 'color' in res: | |
| return res['color'], extrinsics, intrinsics | |
| if 'normal' in res: | |
| return res['normal'], extrinsics, intrinsics | |
| return [], extrinsics, intrinsics | |
| def render_multiview_gpu(sample, resolution=512, nviews=10): | |
| """render_multiview on a black background, frames kept on the GPU as uint8 [N, H, W, 3].""" | |
| cams = [sphere_hammersley_sequence(i, nviews) for i in range(nviews)] | |
| extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics([c[0] for c in cams], [c[1] for c in cams], 2, 40) | |
| renderer = get_renderer(sample, resolution=resolution, bg_color=(0, 0, 0)) | |
| with torch.no_grad(): | |
| frames = torch.stack([ | |
| (renderer.render(sample, extr, intr)['color'].permute(1, 2, 0) * 255).clamp(0, 255).to(torch.uint8) | |
| for extr, intr in zip(extrinsics, intrinsics) | |
| ]) | |
| return frames, extrinsics, intrinsics | |
| def render_snapshot(samples, resolution=512, bg_color=(0, 0, 0), offset=(-16 / 180 * np.pi, 20 / 180 * np.pi), r=10, fov=8, **kwargs): | |
| yaw = [0, np.pi/2, np.pi, 3*np.pi/2] | |
| yaw_offset = offset[0] | |
| yaw = [y + yaw_offset for y in yaw] | |
| pitch = [offset[1] for _ in range(4)] | |
| extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, r, fov) | |
| return render_frames(samples, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs) | |
| class _FrozenGaussian: | |
| def __init__(self, gaussian): | |
| self.get_xyz = gaussian.get_xyz | |
| self.get_opacity = gaussian.get_opacity | |
| self.get_scaling = gaussian.get_scaling | |
| self.get_rotation = gaussian.get_rotation | |
| self.get_features = gaussian.get_features | |
| self.active_sh_degree = gaussian.active_sh_degree | |
| def _turntable_cameras(num_frames, resolution, r, fov, near, far): | |
| from ..renderers.gaussian_render import intrinsics_to_projection | |
| yaws = torch.linspace(0, 2 * 3.1415, num_frames) | |
| pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames)) | |
| extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws.tolist(), pitch.tolist(), r, fov) | |
| views = torch.stack(extrinsics) | |
| centers = torch.inverse(views)[:, :3, 3] | |
| fovs = (2 * torch.atan(0.5 / torch.stack([intr[0, 0] for intr in intrinsics]))).tolist() | |
| cameras = [] | |
| for view, intr, center, fov_rad in zip(views, intrinsics, centers, fovs): | |
| perspective = intrinsics_to_projection(intr, near, far) | |
| cameras.append(edict({ | |
| "image_height": resolution, "image_width": resolution, "FoVx": fov_rad, "FoVy": fov_rad, | |
| "znear": near, "zfar": far, "world_view_transform": view.T.contiguous(), | |
| "projection_matrix": perspective.T.contiguous(), | |
| "full_proj_transform": (perspective @ view).T.contiguous(), "camera_center": center, | |
| })) | |
| return cameras | |
| def render_gaussian_turntable(sample, resolution, num_frames, bg_colors, r=2, fov=40): | |
| """Same frames as render_video(mode='color') per background, kept on the GPU as uint8 [N, H, W, 3].""" | |
| from ..renderers.gaussian_render import render | |
| with torch.no_grad(): | |
| renderer = get_renderer(sample, resolution=resolution) | |
| options = renderer.rendering_options | |
| cameras = _turntable_cameras(num_frames, resolution, r, fov, options.near, options.far) | |
| frozen = _FrozenGaussian(sample) | |
| backgrounds = [torch.tensor(bg, dtype=torch.float32, device="cuda") for bg in bg_colors] | |
| frames = torch.empty(len(bg_colors), num_frames, resolution, resolution, 3, dtype=torch.uint8, device="cuda") | |
| for j, camera in enumerate(cameras): | |
| for b, background in enumerate(backgrounds): | |
| color = render(camera, frozen, renderer.pipe, background, scaling_modifier=renderer.pipe.scale_modifier).render | |
| frames[b, j] = (color.permute(1, 2, 0) * 255).clamp(0, 255).to(torch.uint8) | |
| return frames | |