Upload 21 files

configs/image.yaml

@@ -0,0 +1,69 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 1
+
+### Output
+outdir: logs
+mesh_format: obj
+save_path: ???
+
+### Training
+# guidance loss weights (0 to disable)
+lambda_sd: 0
+lambda_zero123: 1
+# training batch size per iter
+batch_size: 1
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 2
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.5
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+num_pts: 5000
+sh_degree: 0
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+percent_dense: 0.1
+density_start_iter: 100
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.5
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2

configs/text.yaml

@@ -0,0 +1,68 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 1
+
+### Output
+outdir: logs
+mesh_format: obj
+save_path: ???
+
+### Training
+# guidance loss weights (0 to disable)
+lambda_sd: 1
+lambda_zero123: 0
+# training batch size per iter
+batch_size: 1
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 2.5
+# training camera fovy
+fovy: 49.1
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.5
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+num_pts: 1000
+sh_degree: 0
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+percent_dense: 0.1
+density_start_iter: 100
+density_end_iter: 3000
+densification_interval: 50
+opacity_reset_interval: 700
+densify_grad_threshold: 0.01
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2

guidance/sd_utils.py

@@ -0,0 +1,334 @@
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+from diffusers import (
+    AutoencoderKL,
+    UNet2DConditionModel,
+    PNDMScheduler,
+    DDIMScheduler,
+    StableDiffusionPipeline,
+)
+from diffusers.utils.import_utils import is_xformers_available
+
+# suppress partial model loading warning
+logging.set_verbosity_error()
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def seed_everything(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    # torch.backends.cudnn.deterministic = True
+    # torch.backends.cudnn.benchmark = True
+
+
+class StableDiffusion(nn.Module):
+    def __init__(
+        self,
+        device,
+        fp16=True,
+        vram_O=False,
+        sd_version="2.1",
+        hf_key=None,
+        t_range=[0.02, 0.98],
+    ):
+        super().__init__()
+
+        self.device = device
+        self.sd_version = sd_version
+
+        if hf_key is not None:
+            print(f"[INFO] using hugging face custom model key: {hf_key}")
+            model_key = hf_key
+        elif self.sd_version == "2.1":
+            model_key = "stabilityai/stable-diffusion-2-1-base"
+        elif self.sd_version == "2.0":
+            model_key = "stabilityai/stable-diffusion-2-base"
+        elif self.sd_version == "1.5":
+            model_key = "runwayml/stable-diffusion-v1-5"
+        else:
+            raise ValueError(
+                f"Stable-diffusion version {self.sd_version} not supported."
+            )
+
+        self.dtype = torch.float16 if fp16 else torch.float32
+
+        # Create model
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_key, torch_dtype=self.dtype
+        )
+
+        if vram_O:
+            pipe.enable_sequential_cpu_offload()
+            pipe.enable_vae_slicing()
+            pipe.unet.to(memory_format=torch.channels_last)
+            pipe.enable_attention_slicing(1)
+            # pipe.enable_model_cpu_offload()
+        else:
+            pipe.to(device)
+
+        self.vae = pipe.vae
+        self.tokenizer = pipe.tokenizer
+        self.text_encoder = pipe.text_encoder
+        self.unet = pipe.unet
+
+        self.scheduler = DDIMScheduler.from_pretrained(
+            model_key, subfolder="scheduler", torch_dtype=self.dtype
+        )
+
+        del pipe
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.min_step = int(self.num_train_timesteps * t_range[0])
+        self.max_step = int(self.num_train_timesteps * t_range[1])
+        self.alphas = self.scheduler.alphas_cumprod.to(self.device)  # for convenience
+
+        self.embeddings = None
+
+    @torch.no_grad()
+    def get_text_embeds(self, prompts, negative_prompts):
+        pos_embeds = self.encode_text(prompts)  # [1, 77, 768]
+        neg_embeds = self.encode_text(negative_prompts)
+        self.embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)  # [2, 77, 768]
+    
+    def encode_text(self, prompt):
+        # prompt: [str]
+        inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        )
+        embeddings = self.text_encoder(inputs.input_ids.to(self.device))[0]
+        return embeddings
+
+    @torch.no_grad()
+    def refine(self, pred_rgb,
+               guidance_scale=100, steps=50, strength=0.8,
+        ):
+
+        batch_size = pred_rgb.shape[0]
+        pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False)
+        latents = self.encode_imgs(pred_rgb_512.to(self.dtype))
+        # latents = torch.randn((1, 4, 64, 64), device=self.device, dtype=self.dtype)
+
+        self.scheduler.set_timesteps(steps)
+        init_step = int(steps * strength)
+        latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
+
+        for i, t in enumerate(self.scheduler.timesteps[init_step:]):
+    
+            latent_model_input = torch.cat([latents] * 2)
+
+            noise_pred = self.unet(
+                latent_model_input, t, encoder_hidden_states=self.embeddings,
+            ).sample
+
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+            
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        imgs = self.decode_latents(latents) # [1, 3, 512, 512]
+        return imgs
+
+    def train_step(
+        self,
+        pred_rgb,
+        step_ratio=None,
+        guidance_scale=100,
+        as_latent=False,
+    ):
+        
+        batch_size = pred_rgb.shape[0]
+        pred_rgb = pred_rgb.to(self.dtype)
+
+        if as_latent:
+            latents = F.interpolate(pred_rgb, (64, 64), mode="bilinear", align_corners=False) * 2 - 1
+        else:
+            # interp to 512x512 to be fed into vae.
+            pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode="bilinear", align_corners=False)
+            # encode image into latents with vae, requires grad!
+            latents = self.encode_imgs(pred_rgb_512)
+
+        if step_ratio is not None:
+            # dreamtime-like
+            # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
+            t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+        else:
+            t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
+
+        # w(t), sigma_t^2
+        w = (1 - self.alphas[t]).view(batch_size, 1, 1, 1)
+
+        # predict the noise residual with unet, NO grad!
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2)
+            tt = torch.cat([t] * 2)
+
+            noise_pred = self.unet(
+                latent_model_input, tt, encoder_hidden_states=self.embeddings.repeat(batch_size, 1, 1)
+            ).sample
+
+            # perform guidance (high scale from paper!)
+            noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (
+                noise_pred_pos - noise_pred_uncond
+            )
+
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+
+        # seems important to avoid NaN...
+        # grad = grad.clamp(-1, 1)
+
+        target = (latents - grad).detach()
+        loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[0]
+
+        return loss
+
+    @torch.no_grad()
+    def produce_latents(
+        self,
+        height=512,
+        width=512,
+        num_inference_steps=50,
+        guidance_scale=7.5,
+        latents=None,
+    ):
+        if latents is None:
+            latents = torch.randn(
+                (
+                    self.embeddings.shape[0] // 2,
+                    self.unet.in_channels,
+                    height // 8,
+                    width // 8,
+                ),
+                device=self.device,
+            )
+
+        self.scheduler.set_timesteps(num_inference_steps)
+
+        for i, t in enumerate(self.scheduler.timesteps):
+            # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+            latent_model_input = torch.cat([latents] * 2)
+            # predict the noise residual
+            noise_pred = self.unet(
+                latent_model_input, t, encoder_hidden_states=self.embeddings
+            ).sample
+
+            # perform guidance
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (
+                noise_pred_cond - noise_pred_uncond
+            )
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        return latents
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+
+        imgs = self.vae.decode(latents).sample
+        imgs = (imgs / 2 + 0.5).clamp(0, 1)
+
+        return imgs
+
+    def encode_imgs(self, imgs):
+        # imgs: [B, 3, H, W]
+
+        imgs = 2 * imgs - 1
+
+        posterior = self.vae.encode(imgs).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+
+        return latents
+
+    def prompt_to_img(
+        self,
+        prompts,
+        negative_prompts="",
+        height=512,
+        width=512,
+        num_inference_steps=50,
+        guidance_scale=7.5,
+        latents=None,
+    ):
+        if isinstance(prompts, str):
+            prompts = [prompts]
+
+        if isinstance(negative_prompts, str):
+            negative_prompts = [negative_prompts]
+
+        # Prompts -> text embeds
+        self.get_text_embeds(prompts, negative_prompts)
+        
+        # Text embeds -> img latents
+        latents = self.produce_latents(
+            height=height,
+            width=width,
+            latents=latents,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+        )  # [1, 4, 64, 64]
+
+        # Img latents -> imgs
+        imgs = self.decode_latents(latents)  # [1, 3, 512, 512]
+
+        # Img to Numpy
+        imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
+        imgs = (imgs * 255).round().astype("uint8")
+
+        return imgs
+
+
+if __name__ == "__main__":
+    import argparse
+    import matplotlib.pyplot as plt
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("prompt", type=str)
+    parser.add_argument("--negative", default="", type=str)
+    parser.add_argument(
+        "--sd_version",
+        type=str,
+        default="2.1",
+        choices=["1.5", "2.0", "2.1"],
+        help="stable diffusion version",
+    )
+    parser.add_argument(
+        "--hf_key",
+        type=str,
+        default=None,
+        help="hugging face Stable diffusion model key",
+    )
+    parser.add_argument("--fp16", action="store_true", help="use float16 for training")
+    parser.add_argument(
+        "--vram_O", action="store_true", help="optimization for low VRAM usage"
+    )
+    parser.add_argument("-H", type=int, default=512)
+    parser.add_argument("-W", type=int, default=512)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--steps", type=int, default=50)
+    opt = parser.parse_args()
+
+    seed_everything(opt.seed)
+
+    device = torch.device("cuda")
+
+    sd = StableDiffusion(device, opt.fp16, opt.vram_O, opt.sd_version, opt.hf_key)
+
+    imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.H, opt.W, opt.steps)
+
+    # visualize image
+    plt.imshow(imgs[0])
+    plt.show()

guidance/zero123_utils.py

@@ -0,0 +1,226 @@
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+from diffusers import (
+    AutoencoderKL,
+    UNet2DConditionModel,
+    DDIMScheduler,
+    StableDiffusionPipeline,
+)
+import torchvision.transforms.functional as TF
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import sys
+sys.path.append('./')
+
+from zero123 import Zero123Pipeline
+
+
+class Zero123(nn.Module):
+    def __init__(self, device, fp16=True, t_range=[0.02, 0.98]):
+        super().__init__()
+
+        self.device = device
+        self.fp16 = fp16
+        self.dtype = torch.float16 if fp16 else torch.float32
+
+        self.pipe = Zero123Pipeline.from_pretrained(            
+            # "bennyguo/zero123-diffusers",
+            "bennyguo/zero123-xl-diffusers",
+            # './model_cache/zero123_xl',
+            variant="fp16_ema" if self.fp16 else None,
+            torch_dtype=self.dtype,
+        ).to(self.device)
+
+        # for param in self.pipe.parameters():
+        #     param.requires_grad = False
+
+        self.pipe.image_encoder.eval()
+        self.pipe.vae.eval()
+        self.pipe.unet.eval()
+        self.pipe.clip_camera_projection.eval()
+
+        self.vae = self.pipe.vae
+        self.unet = self.pipe.unet
+
+        self.pipe.set_progress_bar_config(disable=True)
+
+        self.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+
+        self.min_step = int(self.num_train_timesteps * t_range[0])
+        self.max_step = int(self.num_train_timesteps * t_range[1])
+        self.alphas = self.scheduler.alphas_cumprod.to(self.device) # for convenience
+
+        self.embeddings = None
+
+    @torch.no_grad()
+    def get_img_embeds(self, x):
+        # x: image tensor in [0, 1]
+        x = F.interpolate(x, (256, 256), mode='bilinear', align_corners=False)
+        x_pil = [TF.to_pil_image(image) for image in x]
+        x_clip = self.pipe.feature_extractor(images=x_pil, return_tensors="pt").pixel_values.to(device=self.device, dtype=self.dtype)
+        c = self.pipe.image_encoder(x_clip).image_embeds
+        v = self.encode_imgs(x.to(self.dtype)) / self.vae.config.scaling_factor
+        self.embeddings = [c, v]
+
+    @torch.no_grad()
+    def refine(self, pred_rgb, polar, azimuth, radius, 
+               guidance_scale=5, steps=50, strength=0.8,
+        ):
+
+        batch_size = pred_rgb.shape[0]
+
+        self.scheduler.set_timesteps(steps)
+
+        if strength == 0:
+            init_step = 0
+            latents = torch.randn((1, 4, 32, 32), device=self.device, dtype=self.dtype)
+        else:
+            init_step = int(steps * strength)
+            pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
+            latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
+            latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
+
+        T = np.stack([np.deg2rad(polar), np.sin(np.deg2rad(azimuth)), np.cos(np.deg2rad(azimuth)), radius], axis=-1)
+        T = torch.from_numpy(T).unsqueeze(1).to(self.dtype).to(self.device) # [8, 1, 4]
+        cc_emb = torch.cat([self.embeddings[0].repeat(batch_size, 1, 1), T], dim=-1)
+        cc_emb = self.pipe.clip_camera_projection(cc_emb)
+        cc_emb = torch.cat([cc_emb, torch.zeros_like(cc_emb)], dim=0)
+
+        vae_emb = self.embeddings[1].repeat(batch_size, 1, 1, 1)
+        vae_emb = torch.cat([vae_emb, torch.zeros_like(vae_emb)], dim=0)
+
+        for i, t in enumerate(self.scheduler.timesteps[init_step:]):
+            
+            x_in = torch.cat([latents] * 2)
+            t_in = torch.cat([t.view(1)] * 2).to(self.device)
+
+            noise_pred = self.unet(
+                torch.cat([x_in, vae_emb], dim=1),
+                t_in.to(self.unet.dtype),
+                encoder_hidden_states=cc_emb,
+            ).sample
+
+            noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+            
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        imgs = self.decode_latents(latents) # [1, 3, 256, 256]
+        return imgs
+    
+    def train_step(self, pred_rgb, polar, azimuth, radius, step_ratio=None, guidance_scale=5, as_latent=False):
+        # pred_rgb: tensor [1, 3, H, W] in [0, 1]
+
+        batch_size = pred_rgb.shape[0]
+
+        if as_latent:
+            latents = F.interpolate(pred_rgb, (32, 32), mode='bilinear', align_corners=False) * 2 - 1
+        else:
+            pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
+            latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
+
+        if step_ratio is not None:
+            # dreamtime-like
+            # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
+            t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+        else:
+            t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
+
+        w = (1 - self.alphas[t]).view(batch_size, 1, 1, 1)
+
+        with torch.no_grad():
+            noise = torch.randn_like(latents)
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+
+            x_in = torch.cat([latents_noisy] * 2)
+            t_in = torch.cat([t] * 2)
+
+            T = np.stack([np.deg2rad(polar), np.sin(np.deg2rad(azimuth)), np.cos(np.deg2rad(azimuth)), radius], axis=-1)
+            T = torch.from_numpy(T).unsqueeze(1).to(self.dtype).to(self.device) # [8, 1, 4]
+            cc_emb = torch.cat([self.embeddings[0].repeat(batch_size, 1, 1), T], dim=-1)
+            cc_emb = self.pipe.clip_camera_projection(cc_emb)
+            cc_emb = torch.cat([cc_emb, torch.zeros_like(cc_emb)], dim=0)
+
+            vae_emb = self.embeddings[1].repeat(batch_size, 1, 1, 1)
+            vae_emb = torch.cat([vae_emb, torch.zeros_like(vae_emb)], dim=0)
+
+            noise_pred = self.unet(
+                torch.cat([x_in, vae_emb], dim=1),
+                t_in.to(self.unet.dtype),
+                encoder_hidden_states=cc_emb,
+            ).sample
+
+        noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+
+        target = (latents - grad).detach()
+        loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum')
+
+        return loss
+    
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+
+        imgs = self.vae.decode(latents).sample
+        imgs = (imgs / 2 + 0.5).clamp(0, 1)
+
+        return imgs
+
+    def encode_imgs(self, imgs, mode=False):
+        # imgs: [B, 3, H, W]
+
+        imgs = 2 * imgs - 1
+
+        posterior = self.vae.encode(imgs).latent_dist
+        if mode:
+            latents = posterior.mode()
+        else:
+            latents = posterior.sample() 
+        latents = latents * self.vae.config.scaling_factor
+
+        return latents
+    
+    
+if __name__ == '__main__':
+    import cv2
+    import argparse
+    import numpy as np
+    import matplotlib.pyplot as plt
+    
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('input', type=str)
+    parser.add_argument('--polar', type=float, default=0, help='delta polar angle in [-90, 90]')
+    parser.add_argument('--azimuth', type=float, default=0, help='delta azimuth angle in [-180, 180]')
+    parser.add_argument('--radius', type=float, default=0, help='delta camera radius multiplier in [-0.5, 0.5]')
+
+    opt = parser.parse_args()
+
+    device = torch.device('cuda')
+
+    print(f'[INFO] loading image from {opt.input} ...')
+    image = cv2.imread(opt.input, cv2.IMREAD_UNCHANGED)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    image = cv2.resize(image, (256, 256), interpolation=cv2.INTER_AREA)
+    image = image.astype(np.float32) / 255.0
+    image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).contiguous().to(device)
+
+    print(f'[INFO] loading model ...')
+    zero123 = Zero123(device)
+
+    print(f'[INFO] running model ...')
+    zero123.get_img_embeds(image)
+
+    while True:
+        outputs = zero123.refine(image, polar=[opt.polar], azimuth=[opt.azimuth], radius=[opt.radius], strength=0)
+        plt.imshow(outputs.float().cpu().numpy().transpose(0, 2, 3, 1)[0])
+        plt.show()

scripts/convert_obj_to_video.py

@@ -0,0 +1,20 @@
+import os
+import glob
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--dir', default='logs', type=str, help='Directory where obj files are stored')
+parser.add_argument('--out', default='videos', type=str, help='Directory where videos will be saved')
+args = parser.parse_args()
+
+out = args.out
+os.makedirs(out, exist_ok=True)
+
+files = glob.glob(f'{args.dir}/*.obj')
+for f in files:
+    name = os.path.basename(f)
+    # first stage model, ignore
+    if name.endswith('_mesh.obj'): 
+        continue
+    print(f'[INFO] process {name}')
+    os.system(f"python -m kiui.render {f} --save_video {os.path.join(out, name.replace('.obj', '.mp4'))} ")

scripts/run.sh

@@ -0,0 +1,5 @@
+export CUDA_VISIBLE_DEVICES=5
+
+python main.py --config configs/image.yaml input=data/anya_rgba.png save_path=anya
+python main2.py --config configs/image.yaml input=data/anya_rgba.png save_path=anya
+python -m kiui.render logs/anya.obj --save_video videos/anya.mp4 --wogui

scripts/run_sd.sh

@@ -0,0 +1,31 @@
+export CUDA_VISIBLE_DEVICES=6
+
+# easy samples
+python main.py --config configs/text.yaml prompt="a photo of an icecream" save_path=icecream 
+python main2.py --config configs/text.yaml prompt="a photo of an icecream" save_path=icecream 
+python main.py --config configs/text.yaml prompt="a ripe strawberry" save_path=strawberry 
+python main2.py --config configs/text.yaml prompt="a ripe strawberry" save_path=strawberry 
+python main.py --config configs/text.yaml prompt="a blue tulip" save_path=tulip 
+python main2.py --config configs/text.yaml prompt="a blue tulip" save_path=tulip 
+
+python main.py --config configs/text.yaml prompt="a golden goblet" save_path=goblet 
+python main2.py --config configs/text.yaml prompt="a golden goblet" save_path=goblet 
+python main.py --config configs/text.yaml prompt="a photo of a hamburger" save_path=hamburger 
+python main2.py --config configs/text.yaml prompt="a photo of a hamburger" save_path=hamburger 
+python main.py --config configs/text.yaml prompt="a delicious croissant" save_path=croissant 
+python main2.py --config configs/text.yaml prompt="a delicious croissant" save_path=croissant 
+
+# hard samples
+python main.py --config configs/text.yaml prompt="a baby bunny sitting on top of a stack of pancake" save_path=bunny_pancake 
+python main2.py --config configs/text.yaml prompt="a baby bunny sitting on top of a stack of pancake" save_path=bunny_pancake 
+python main.py --config configs/text.yaml prompt="a typewriter" save_path=typewriter 
+python main2.py --config configs/text.yaml prompt="a typewriter" save_path=typewriter 
+python main.py --config configs/text.yaml prompt="a pineapple" save_path=pineapple 
+python main2.py --config configs/text.yaml prompt="a pineapple" save_path=pineapple 
+
+python main.py --config configs/text.yaml prompt="a model of a house in Tudor style" save_path=tudor_house 
+python main2.py --config configs/text.yaml prompt="a model of a house in Tudor style" save_path=tudor_house 
+python main.py --config configs/text.yaml prompt="a lionfish" save_path=lionfish 
+python main2.py --config configs/text.yaml prompt="a lionfish" save_path=lionfish 
+python main.py --config configs/text.yaml prompt="a bunch of yellow rose, highly detailed" save_path=rose 
+python main2.py --config configs/text.yaml prompt="a bunch of yellow rose, highly detailed" save_path=rose 

scripts/runall.py

@@ -0,0 +1,48 @@
+import os
+import glob
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--dir', default='data', type=str, help='Directory where processed images are stored')
+parser.add_argument('--out', default='logs', type=str, help='Directory where obj files will be saved')
+parser.add_argument('--video-out', default='videos', type=str, help='Directory where videos will be saved')
+parser.add_argument('--gpu', default=0, type=int, help='ID of GPU to use')
+parser.add_argument('--elevation', default=0, type=int, help='Elevation angle of view in degrees')
+parser.add_argument('--config', default='configs', type=str, help='Path to config directory, which contains image.yaml')
+args = parser.parse_args()
+
+files = glob.glob(f'{args.dir}/*_rgba.png')
+configs_dir = args.config
+
+# check if image.yaml exists
+if not os.path.exists(os.path.join(configs_dir, 'image.yaml')):
+    raise FileNotFoundError(
+        f'image.yaml not found in {configs_dir} directory. Please check if the directory is correct.'
+    )
+
+# create output directories if not exists
+out_dir = args.out
+os.makedirs(out_dir, exist_ok=True)
+video_dir = args.video_out
+os.makedirs(video_dir, exist_ok=True)
+
+
+for file in files:
+    name = os.path.basename(file).replace("_rgba.png", "")
+    print(f'======== processing {name} ========')
+    # first stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py '
+              f'--config {configs_dir}/image.yaml '
+              f'input={file} '
+              f'save_path={name} elevation={args.elevation}')
+    # second stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py '
+              f'--config {configs_dir}/image.yaml '
+              f'input={file} '
+              f'save_path={name} elevation={args.elevation}')
+    # export video
+    mesh_path = os.path.join(out_dir, f'{name}.obj')
+    os.system(f'python -m kiui.render {mesh_path} '
+              f'--save_video {video_dir}/{name}.mp4 '
+              f'--wogui '
+              f'--elevation {args.elevation}')

scripts/runall_sd.py

@@ -0,0 +1,45 @@
+import os
+import glob
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu', default=0, type=int)
+args = parser.parse_args()
+
+prompts = [
+    ('strawberry', 'a ripe strawberry'),
+    ('cactus_pot', 'a small saguaro cactus planted in a clay pot'),
+    ('hamburger', 'a delicious hamburger'),
+    ('icecream', 'an icecream'),
+    ('tulip', 'a blue tulip'),
+    ('pineapple', 'a ripe pineapple'),
+    ('goblet', 'a golden goblet'),
+    # ('squitopus', 'a squirrel-octopus hybrid'),
+    # ('astronaut', 'Michelangelo style statue of an astronaut'),
+    # ('teddy_bear', 'a teddy bear'),
+    # ('corgi_nurse', 'a plush toy of a corgi nurse'),
+    # ('teapot', 'a blue and white porcelain teapot'),
+    # ('skull', "a human skull"),
+    # ('penguin', 'a penguin'),
+    # ('campfire', 'a campfire'),
+    # ('donut', 'a donut with pink icing'),
+    # ('cupcake', 'a birthday cupcake'),
+    # ('pie', 'shepherds pie'),
+    # ('cone', 'a traffic cone'),
+    # ('schoolbus', 'a schoolbus'),
+    # ('avocado_chair', 'a chair that looks like an avocado'),
+    # ('glasses', 'a pair of sunglasses')
+    # ('potion', 'a bottle of green potion'),
+    # ('chalice', 'a delicate chalice'),
+]
+
+for name, prompt in prompts:
+    print(f'======== processing {name} ========')
+    # first stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py --config configs/text.yaml prompt="{prompt}" save_path={name}')
+    # second stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py --config configs/text.yaml  prompt="{prompt}" save_path={name}')
+    # export video
+    mesh_path = os.path.join('logs', f'{name}.obj')
+    os.makedirs('videos', exist_ok=True)
+    os.system(f'python -m kiui.render {mesh_path} --save_video videos/{name}.mp4 --wogui')

simple-knn/ext.cpp

@@ -0,0 +1,17 @@
+/*
+ * Copyright (C) 2023, Inria
+ * GRAPHDECO research group, https://team.inria.fr/graphdeco
+ * All rights reserved.
+ *
+ * This software is free for non-commercial, research and evaluation use 
+ * under the terms of the LICENSE.md file.
+ *
+ * For inquiries contact  [email protected]
+ */
+
+#include <torch/extension.h>
+#include "spatial.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("distCUDA2", &distCUDA2);
+}

simple-knn/setup.py

@@ -0,0 +1,35 @@
+#
+# Copyright (C) 2023, Inria
+# GRAPHDECO research group, https://team.inria.fr/graphdeco
+# All rights reserved.
+#
+# This software is free for non-commercial, research and evaluation use 
+# under the terms of the LICENSE.md file.
+#
+# For inquiries contact  [email protected]
+#
+
+from setuptools import setup
+from torch.utils.cpp_extension import CUDAExtension, BuildExtension
+import os
+
+cxx_compiler_flags = []
+
+if os.name == 'nt':
+    cxx_compiler_flags.append("/wd4624")
+
+setup(
+    name="simple_knn",
+    ext_modules=[
+        CUDAExtension(
+            name="simple_knn._C",
+            sources=[
+            "spatial.cu", 
+            "simple_knn.cu",
+            "ext.cpp"],
+            extra_compile_args={"nvcc": [], "cxx": cxx_compiler_flags})
+        ],
+    cmdclass={
+        'build_ext': BuildExtension
+    }
+)

simple-knn/simple_knn.cu

@@ -0,0 +1,221 @@
+/*
+ * Copyright (C) 2023, Inria
+ * GRAPHDECO research group, https://team.inria.fr/graphdeco
+ * All rights reserved.
+ *
+ * This software is free for non-commercial, research and evaluation use 
+ * under the terms of the LICENSE.md file.
+ *
+ * For inquiries contact  [email protected]
+ */
+
+#define BOX_SIZE 1024
+
+#include "cuda_runtime.h"
+#include "device_launch_parameters.h"
+#include "simple_knn.h"
+#include <cub/cub.cuh>
+#include <cub/device/device_radix_sort.cuh>
+#include <vector>
+#include <cuda_runtime_api.h>
+#include <thrust/device_vector.h>
+#include <thrust/sequence.h>
+#define __CUDACC__
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace cg = cooperative_groups;
+
+struct CustomMin
+{
+	__device__ __forceinline__
+		float3 operator()(const float3& a, const float3& b) const {
+		return { min(a.x, b.x), min(a.y, b.y), min(a.z, b.z) };
+	}
+};
+
+struct CustomMax
+{
+	__device__ __forceinline__
+		float3 operator()(const float3& a, const float3& b) const {
+		return { max(a.x, b.x), max(a.y, b.y), max(a.z, b.z) };
+	}
+};
+
+__host__ __device__ uint32_t prepMorton(uint32_t x)
+{
+	x = (x | (x << 16)) & 0x030000FF;
+	x = (x | (x << 8)) & 0x0300F00F;
+	x = (x | (x << 4)) & 0x030C30C3;
+	x = (x | (x << 2)) & 0x09249249;
+	return x;
+}
+
+__host__ __device__ uint32_t coord2Morton(float3 coord, float3 minn, float3 maxx)
+{
+	uint32_t x = prepMorton(((coord.x - minn.x) / (maxx.x - minn.x)) * ((1 << 10) - 1));
+	uint32_t y = prepMorton(((coord.y - minn.y) / (maxx.y - minn.y)) * ((1 << 10) - 1));
+	uint32_t z = prepMorton(((coord.z - minn.z) / (maxx.z - minn.z)) * ((1 << 10) - 1));
+
+	return x | (y << 1) | (z << 2);
+}
+
+__global__ void coord2Morton(int P, const float3* points, float3 minn, float3 maxx, uint32_t* codes)
+{
+	auto idx = cg::this_grid().thread_rank();
+	if (idx >= P)
+		return;
+
+	codes[idx] = coord2Morton(points[idx], minn, maxx);
+}
+
+struct MinMax
+{
+	float3 minn;
+	float3 maxx;
+};
+
+__global__ void boxMinMax(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes)
+{
+	auto idx = cg::this_grid().thread_rank();
+
+	MinMax me;
+	if (idx < P)
+	{
+		me.minn = points[indices[idx]];
+		me.maxx = points[indices[idx]];
+	}
+	else
+	{
+		me.minn = { FLT_MAX, FLT_MAX, FLT_MAX };
+		me.maxx = { -FLT_MAX,-FLT_MAX,-FLT_MAX };
+	}
+
+	__shared__ MinMax redResult[BOX_SIZE];
+
+	for (int off = BOX_SIZE / 2; off >= 1; off /= 2)
+	{
+		if (threadIdx.x < 2 * off)
+			redResult[threadIdx.x] = me;
+		__syncthreads();
+
+		if (threadIdx.x < off)
+		{
+			MinMax other = redResult[threadIdx.x + off];
+			me.minn.x = min(me.minn.x, other.minn.x);
+			me.minn.y = min(me.minn.y, other.minn.y);
+			me.minn.z = min(me.minn.z, other.minn.z);
+			me.maxx.x = max(me.maxx.x, other.maxx.x);
+			me.maxx.y = max(me.maxx.y, other.maxx.y);
+			me.maxx.z = max(me.maxx.z, other.maxx.z);
+		}
+		__syncthreads();
+	}
+
+	if (threadIdx.x == 0)
+		boxes[blockIdx.x] = me;
+}
+
+__device__ __host__ float distBoxPoint(const MinMax& box, const float3& p)
+{
+	float3 diff = { 0, 0, 0 };
+	if (p.x < box.minn.x || p.x > box.maxx.x)
+		diff.x = min(abs(p.x - box.minn.x), abs(p.x - box.maxx.x));
+	if (p.y < box.minn.y || p.y > box.maxx.y)
+		diff.y = min(abs(p.y - box.minn.y), abs(p.y - box.maxx.y));
+	if (p.z < box.minn.z || p.z > box.maxx.z)
+		diff.z = min(abs(p.z - box.minn.z), abs(p.z - box.maxx.z));
+	return diff.x * diff.x + diff.y * diff.y + diff.z * diff.z;
+}
+
+template<int K>
+__device__ void updateKBest(const float3& ref, const float3& point, float* knn)
+{
+	float3 d = { point.x - ref.x, point.y - ref.y, point.z - ref.z };
+	float dist = d.x * d.x + d.y * d.y + d.z * d.z;
+	for (int j = 0; j < K; j++)
+	{
+		if (knn[j] > dist)
+		{
+			float t = knn[j];
+			knn[j] = dist;
+			dist = t;
+		}
+	}
+}
+
+__global__ void boxMeanDist(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes, float* dists)
+{
+	int idx = cg::this_grid().thread_rank();
+	if (idx >= P)
+		return;
+
+	float3 point = points[indices[idx]];
+	float best[3] = { FLT_MAX, FLT_MAX, FLT_MAX };
+
+	for (int i = max(0, idx - 3); i <= min(P - 1, idx + 3); i++)
+	{
+		if (i == idx)
+			continue;
+		updateKBest<3>(point, points[indices[i]], best);
+	}
+
+	float reject = best[2];
+	best[0] = FLT_MAX;
+	best[1] = FLT_MAX;
+	best[2] = FLT_MAX;
+
+	for (int b = 0; b < (P + BOX_SIZE - 1) / BOX_SIZE; b++)
+	{
+		MinMax box = boxes[b];
+		float dist = distBoxPoint(box, point);
+		if (dist > reject || dist > best[2])
+			continue;
+
+		for (int i = b * BOX_SIZE; i < min(P, (b + 1) * BOX_SIZE); i++)
+		{
+			if (i == idx)
+				continue;
+			updateKBest<3>(point, points[indices[i]], best);
+		}
+	}
+	dists[indices[idx]] = (best[0] + best[1] + best[2]) / 3.0f;
+}
+
+void SimpleKNN::knn(int P, float3* points, float* meanDists)
+{
+	float3* result;
+	cudaMalloc(&result, sizeof(float3));
+	size_t temp_storage_bytes;
+
+	float3 init = { 0, 0, 0 }, minn, maxx;
+
+	cub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, points, result, P, CustomMin(), init);
+	thrust::device_vector<char> temp_storage(temp_storage_bytes);
+
+	cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMin(), init);
+	cudaMemcpy(&minn, result, sizeof(float3), cudaMemcpyDeviceToHost);
+
+	cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMax(), init);
+	cudaMemcpy(&maxx, result, sizeof(float3), cudaMemcpyDeviceToHost);
+
+	thrust::device_vector<uint32_t> morton(P);
+	thrust::device_vector<uint32_t> morton_sorted(P);
+	coord2Morton << <(P + 255) / 256, 256 >> > (P, points, minn, maxx, morton.data().get());
+
+	thrust::device_vector<uint32_t> indices(P);
+	thrust::sequence(indices.begin(), indices.end());
+	thrust::device_vector<uint32_t> indices_sorted(P);
+
+	cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P);
+	temp_storage.resize(temp_storage_bytes);
+
+	cub::DeviceRadixSort::SortPairs(temp_storage.data().get(), temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P);
+
+	uint32_t num_boxes = (P + BOX_SIZE - 1) / BOX_SIZE;
+	thrust::device_vector<MinMax> boxes(num_boxes);
+	boxMinMax << <num_boxes, BOX_SIZE >> > (P, points, indices_sorted.data().get(), boxes.data().get());
+	boxMeanDist << <num_boxes, BOX_SIZE >> > (P, points, indices_sorted.data().get(), boxes.data().get(), meanDists);
+
+	cudaFree(result);
+}

simple-knn/simple_knn.h

@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2023, Inria
+ * GRAPHDECO research group, https://team.inria.fr/graphdeco
+ * All rights reserved.
+ *
+ * This software is free for non-commercial, research and evaluation use 
+ * under the terms of the LICENSE.md file.
+ *
+ * For inquiries contact  [email protected]
+ */
+
+#ifndef SIMPLEKNN_H_INCLUDED
+#define SIMPLEKNN_H_INCLUDED
+
+class SimpleKNN
+{
+public:
+	static void knn(int P, float3* points, float* meanDists);
+};
+
+#endif

simple-knn/spatial.cu

@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2023, Inria
+ * GRAPHDECO research group, https://team.inria.fr/graphdeco
+ * All rights reserved.
+ *
+ * This software is free for non-commercial, research and evaluation use 
+ * under the terms of the LICENSE.md file.
+ *
+ * For inquiries contact  [email protected]
+ */
+
+#include "spatial.h"
+#include "simple_knn.h"
+
+torch::Tensor
+distCUDA2(const torch::Tensor& points)
+{
+  const int P = points.size(0);
+
+  auto float_opts = points.options().dtype(torch::kFloat32);
+  torch::Tensor means = torch::full({P}, 0.0, float_opts);
+  
+  SimpleKNN::knn(P, (float3*)points.contiguous().data<float>(), means.contiguous().data<float>());
+
+  return means;
+}

simple-knn/spatial.h

@@ -0,0 +1,14 @@
+/*
+ * Copyright (C) 2023, Inria
+ * GRAPHDECO research group, https://team.inria.fr/graphdeco
+ * All rights reserved.
+ *
+ * This software is free for non-commercial, research and evaluation use 
+ * under the terms of the LICENSE.md file.
+ *
+ * For inquiries contact  [email protected]
+ */
+
+#include <torch/extension.h>
+
+torch::Tensor distCUDA2(const torch::Tensor& points);