Upload libs/sample.py with huggingface_hub

libs/sample.py

@@ -1,380 +1,380 @@
-import numpy as np
-import torch
-from imagedream.camera_utils import get_camera_for_index
-from imagedream.ldm.util import set_seed, add_random_background
-from libs.base_utils import do_resize_content
-from imagedream.ldm.models.diffusion.ddim import DDIMSampler
-from torchvision import transforms as T
-
-
-class ImageDreamDiffusion:
-    def __init__(
-        self,
-        model,
-        device,
-        dtype,
-        mode,
-        num_frames,
-        camera_views,
-        ref_position,
-        random_background=False,
-        offset_noise=False,
-        resize_rate=1,
-        image_size=256,
-        seed=1234,
-    ) -> None:
-        assert mode in ["pixel", "local"]
-        size = image_size
-        self.seed = seed
-        batch_size = max(4, num_frames)
-
-        neg_texts = "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated,  obscure, unnatural colors, poor lighting, dull, and unclear."
-        uc = model.get_learned_conditioning([neg_texts]).to(device)
-        sampler = DDIMSampler(model)
-
-        # pre-compute camera matrices
-        camera = [get_camera_for_index(i).squeeze() for i in camera_views]
-        camera[ref_position] = torch.zeros_like(camera[ref_position])  # set ref camera to zero
-        camera = torch.stack(camera)
-        camera = camera.repeat(batch_size // num_frames, 1).to(device)
-
-        self.image_transform = T.Compose(
-            [
-                T.Resize((size, size)),
-                T.ToTensor(),
-                T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-            ]
-        )
-        self.dtype = dtype
-        self.ref_position = ref_position
-        self.mode = mode
-        self.random_background = random_background
-        self.resize_rate = resize_rate
-        self.num_frames = num_frames
-        self.size = size
-        self.device = device
-        self.batch_size = batch_size
-        self.model = model
-        self.sampler = sampler
-        self.uc = uc
-        self.camera = camera
-        self.offset_noise = offset_noise
-
-    @staticmethod
-    def i2i(
-        model,
-        image_size,
-        prompt,
-        uc,
-        sampler,
-        ip=None,
-        step=20,
-        scale=5.0,
-        batch_size=8,
-        ddim_eta=0.0,
-        dtype=torch.float32,
-        device="cuda",
-        camera=None,
-        num_frames=4,
-        pixel_control=False,
-        transform=None,
-        offset_noise=False,
-    ):
-        """ The function supports additional image prompt.
-        Args:
-            model (_type_): the image dream model
-            image_size (_type_): size of diffusion output (standard 256)
-            prompt (_type_): text prompt for the image (prompt in type str)
-            uc (_type_): unconditional vector (tensor in shape [1, 77, 1024])
-            sampler (_type_): imagedream.ldm.models.diffusion.ddim.DDIMSampler
-            ip (Image, optional): the image prompt. Defaults to None.
-            step (int, optional): _description_. Defaults to 20.
-            scale (float, optional): _description_. Defaults to 7.5.
-            batch_size (int, optional): _description_. Defaults to 8.
-            ddim_eta (float, optional): _description_. Defaults to 0.0.
-            dtype (_type_, optional): _description_. Defaults to torch.float32.
-            device (str, optional): _description_. Defaults to "cuda".
-            camera (_type_, optional): camera info in tensor, shape: torch.Size([5, 16]) mean: 0.11, std: 0.49, min: -1.00, max: 1.00
-            num_frames (int, optional): _num of frames (views) to generate
-            pixel_control: whether to use pixel conditioning. Defaults to False, True when using pixel mode
-            transform: Compose(
-                Resize(size=(256, 256), interpolation=bilinear, max_size=None, antialias=warn)
-                ToTensor()
-                Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
-            )
-        """
-        ip_raw = ip
-        if type(prompt) != list:
-            prompt = [prompt]
-        with torch.no_grad(), torch.autocast(device_type=torch.device(device).type, dtype=dtype):
-            c = model.get_learned_conditioning(prompt).to(
-                device
-            )  # shape: torch.Size([1, 77, 1024]) mean: -0.17, std: 1.02, min: -7.50, max: 13.05
-            c_ = {"context": c.repeat(batch_size, 1, 1)}  # batch_size
-            uc_ = {"context": uc.repeat(batch_size, 1, 1)}
-
-            if camera is not None:
-                c_["camera"] = uc_["camera"] = (
-                    camera  # shape: torch.Size([5, 16]) mean: 0.11, std: 0.49, min: -1.00, max: 1.00
-                )
-                c_["num_frames"] = uc_["num_frames"] = num_frames
-
-            if ip is not None:
-                ip_embed = model.get_learned_image_conditioning(ip).to(
-                    device
-                )  # shape: torch.Size([1, 257, 1280]) mean: 0.06, std: 0.53, min: -6.83, max: 11.12
-                ip_ = ip_embed.repeat(batch_size, 1, 1)
-                c_["ip"] = ip_
-                uc_["ip"] = torch.zeros_like(ip_)
-
-            if pixel_control:
-                assert camera is not None
-                ip = transform(ip).to(
-                    device
-                )  # shape: torch.Size([3, 256, 256]) mean: 0.33, std: 0.37, min: -1.00, max: 1.00
-                ip_img = model.get_first_stage_encoding(
-                    model.encode_first_stage(ip[None, :, :, :])
-                )  # shape: torch.Size([1, 4, 32, 32]) mean: 0.23, std: 0.77, min: -4.42, max: 3.55
-                c_["ip_img"] = ip_img
-                uc_["ip_img"] = torch.zeros_like(ip_img)
-
-            shape = [4, image_size // 8, image_size // 8]  # [4, 32, 32]
-            if offset_noise:
-                ref = transform(ip_raw).to(device)
-                ref_latent = model.get_first_stage_encoding(model.encode_first_stage(ref[None, :, :, :]))
-                ref_mean = ref_latent.mean(dim=(-1, -2), keepdim=True)
-                time_steps = torch.randint(model.num_timesteps - 1, model.num_timesteps, (batch_size,), device=device)
-                x_T = model.q_sample(torch.ones([batch_size] + shape, device=device) * ref_mean, time_steps)
-
-            samples_ddim, _ = (
-                sampler.sample(  # shape: torch.Size([5, 4, 32, 32]) mean: 0.29, std: 0.85, min: -3.38, max: 4.43
-                    S=step,
-                    conditioning=c_,
-                    batch_size=batch_size,
-                    shape=shape,
-                    verbose=False,
-                    unconditional_guidance_scale=scale,
-                    unconditional_conditioning=uc_,
-                    eta=ddim_eta,
-                    x_T=x_T if offset_noise else None,
-                )
-            )
-
-            x_sample = model.decode_first_stage(samples_ddim)
-            x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
-            x_sample = 255.0 * x_sample.permute(0, 2, 3, 1).cpu().numpy()
-
-        return list(x_sample.astype(np.uint8))
-
-    def diffuse(self, t, ip, n_test=2):
-        set_seed(self.seed)
-        ip = do_resize_content(ip, self.resize_rate)
-        if self.random_background:
-            ip = add_random_background(ip)
-
-        images = []
-        for _ in range(n_test):
-            img = self.i2i(
-                self.model,
-                self.size,
-                t,
-                self.uc,
-                self.sampler,
-                ip=ip,
-                step=50,
-                scale=5,
-                batch_size=self.batch_size,
-                ddim_eta=0.0,
-                dtype=self.dtype,
-                device=self.device,
-                camera=self.camera,
-                num_frames=self.num_frames,
-                pixel_control=(self.mode == "pixel"),
-                transform=self.image_transform,
-                offset_noise=self.offset_noise,
-            )
-            img = np.concatenate(img, 1)
-            img = np.concatenate((img, ip.resize((self.size, self.size))), axis=1)
-            images.append(img)
-        set_seed()  # unset random and numpy seed
-        return images
-
-
-class ImageDreamDiffusionStage2:
-    def __init__(
-        self,
-        model,
-        device,
-        dtype,
-        num_frames,
-        camera_views,
-        ref_position,
-        random_background=False,
-        offset_noise=False,
-        resize_rate=1,
-        mode="pixel",
-        image_size=256,
-        seed=1234,
-    ) -> None:
-        assert mode in ["pixel", "local"]
-
-        size = image_size
-        self.seed = seed
-        batch_size = max(4, num_frames)
-
-        neg_texts = "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated,  obscure, unnatural colors, poor lighting, dull, and unclear."
-        uc = model.get_learned_conditioning([neg_texts]).to(device)
-        sampler = DDIMSampler(model)
-
-        # pre-compute camera matrices
-        camera = [get_camera_for_index(i).squeeze() for i in camera_views]
-        if ref_position is not None:
-            camera[ref_position] = torch.zeros_like(camera[ref_position])  # set ref camera to zero
-        camera = torch.stack(camera)
-        camera = camera.repeat(batch_size // num_frames, 1).to(device)
-
-        self.image_transform = T.Compose(
-            [
-                T.Resize((size, size)),
-                T.ToTensor(),
-                T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-            ]
-        )
-
-        self.dtype = dtype
-        self.mode = mode
-        self.ref_position = ref_position
-        self.random_background = random_background
-        self.resize_rate = resize_rate
-        self.num_frames = num_frames
-        self.size = size
-        self.device = device
-        self.batch_size = batch_size
-        self.model = model
-        self.sampler = sampler
-        self.uc = uc
-        self.camera = camera
-        self.offset_noise = offset_noise
-
-    @staticmethod
-    def i2iStage2(
-        model,
-        image_size,
-        prompt,
-        uc,
-        sampler,
-        pixel_images,
-        ip=None,
-        step=20,
-        scale=5.0,
-        batch_size=8,
-        ddim_eta=0.0,
-        dtype=torch.float32,
-        device="cuda",
-        camera=None,
-        num_frames=4,
-        pixel_control=False,
-        transform=None,
-        offset_noise=False,
-    ):
-        ip_raw = ip
-        if type(prompt) != list:
-            prompt = [prompt]
-        with torch.no_grad(), torch.autocast(device_type=torch.device(device).type, dtype=dtype):
-            c = model.get_learned_conditioning(prompt).to(
-                device
-            )  # shape: torch.Size([1, 77, 1024]) mean: -0.17, std: 1.02, min: -7.50, max: 13.05
-            c_ = {"context": c.repeat(batch_size, 1, 1)}  # batch_size
-            uc_ = {"context": uc.repeat(batch_size, 1, 1)}
-
-            if camera is not None:
-                c_["camera"] = uc_["camera"] = (
-                    camera  # shape: torch.Size([5, 16]) mean: 0.11, std: 0.49, min: -1.00, max: 1.00
-                )
-                c_["num_frames"] = uc_["num_frames"] = num_frames
-
-            if ip is not None:
-                ip_embed = model.get_learned_image_conditioning(ip).to(
-                    device
-                )  # shape: torch.Size([1, 257, 1280]) mean: 0.06, std: 0.53, min: -6.83, max: 11.12
-                ip_ = ip_embed.repeat(batch_size, 1, 1)
-                c_["ip"] = ip_
-                uc_["ip"] = torch.zeros_like(ip_)
-
-            if pixel_control:
-                assert camera is not None
-                
-            transed_pixel_images = torch.stack([transform(i).to(device) for i in pixel_images])
-            latent_pixel_images = model.get_first_stage_encoding(model.encode_first_stage(transed_pixel_images))
-
-            c_["pixel_images"] = latent_pixel_images
-            uc_["pixel_images"] = torch.zeros_like(latent_pixel_images)
-
-            shape = [4, image_size // 8, image_size // 8]  # [4, 32, 32]
-            if offset_noise:
-                ref = transform(ip_raw).to(device)
-                ref_latent = model.get_first_stage_encoding(model.encode_first_stage(ref[None, :, :, :]))
-                ref_mean = ref_latent.mean(dim=(-1, -2), keepdim=True)
-                time_steps = torch.randint(model.num_timesteps - 1, model.num_timesteps, (batch_size,), device=device)
-                x_T = model.q_sample(torch.ones([batch_size] + shape, device=device) * ref_mean, time_steps)
-
-            samples_ddim, _ = (
-                sampler.sample(  # shape: torch.Size([5, 4, 32, 32]) mean: 0.29, std: 0.85, min: -3.38, max: 4.43
-                    S=step,
-                    conditioning=c_,
-                    batch_size=batch_size,
-                    shape=shape,
-                    verbose=False,
-                    unconditional_guidance_scale=scale,
-                    unconditional_conditioning=uc_,
-                    eta=ddim_eta,
-                    x_T=x_T if offset_noise else None,
-                )
-            )
-            x_sample = model.decode_first_stage(samples_ddim)
-            x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
-            x_sample = 255.0 * x_sample.permute(0, 2, 3, 1).cpu().numpy()
-
-        return list(x_sample.astype(np.uint8))
-
-    @torch.no_grad()
-    def diffuse(self, t, ip, pixel_images, n_test=2):
-        set_seed(self.seed)
-        ip = do_resize_content(ip, self.resize_rate)
-        pixel_images = [do_resize_content(i, self.resize_rate) for i in pixel_images]
-
-        if self.random_background:
-            bg_color = np.random.rand() * 255
-            ip = add_random_background(ip, bg_color)
-            pixel_images = [add_random_background(i, bg_color) for i in pixel_images]
-
-        images = []
-        for _ in range(n_test):
-            img = self.i2iStage2(
-                self.model,
-                self.size,
-                t,
-                self.uc,
-                self.sampler,
-                pixel_images=pixel_images,
-                ip=ip,
-                step=50,
-                scale=5,
-                batch_size=self.batch_size,
-                ddim_eta=0.0,
-                dtype=self.dtype,
-                device=self.device,
-                camera=self.camera,
-                num_frames=self.num_frames,
-                pixel_control=(self.mode == "pixel"),
-                transform=self.image_transform,
-                offset_noise=self.offset_noise,
-            )
-            img = np.concatenate(img, 1)
-            img = np.concatenate(
-                (img, ip.resize((self.size, self.size)), *[i.resize((self.size, self.size)) for i in pixel_images]),
-                axis=1,
-            )
-            images.append(img)
-        set_seed()  # unset random and numpy seed
-        return images
+import numpy as np
+import torch
+from imagedream.camera_utils import get_camera_for_index
+from imagedream.ldm.util import set_seed, add_random_background
+from libs.base_utils import do_resize_content
+from imagedream.ldm.models.diffusion.ddim import DDIMSampler
+from torchvision import transforms as T
+
+
+class ImageDreamDiffusion:
+    def __init__(
+        self,
+        model,
+        device,
+        dtype,
+        mode,
+        num_frames,
+        camera_views,
+        ref_position,
+        random_background=False,
+        offset_noise=False,
+        resize_rate=1,
+        image_size=256,
+        seed=1234,
+    ) -> None:
+        assert mode in ["pixel", "local"]
+        size = image_size
+        self.seed = seed
+        batch_size = max(4, num_frames)
+
+        neg_texts = "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated,  obscure, unnatural colors, poor lighting, dull, and unclear."
+        uc = model.get_learned_conditioning([neg_texts]).to(device)
+        sampler = DDIMSampler(model)
+
+        # pre-compute camera matrices
+        camera = [get_camera_for_index(i).squeeze() for i in camera_views]
+        camera[ref_position] = torch.zeros_like(camera[ref_position])  # set ref camera to zero
+        camera = torch.stack(camera)
+        camera = camera.repeat(batch_size // num_frames, 1).to(device)
+
+        self.image_transform = T.Compose(
+            [
+                T.Resize((size, size)),
+                T.ToTensor(),
+                T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+            ]
+        )
+        self.dtype = dtype
+        self.ref_position = ref_position
+        self.mode = mode
+        self.random_background = random_background
+        self.resize_rate = resize_rate
+        self.num_frames = num_frames
+        self.size = size
+        self.device = device
+        self.batch_size = batch_size
+        self.model = model
+        self.sampler = sampler
+        self.uc = uc
+        self.camera = camera
+        self.offset_noise = offset_noise
+
+    @staticmethod
+    def i2i(
+        model,
+        image_size,
+        prompt,
+        uc,
+        sampler,
+        ip=None,
+        step=20,
+        scale=5.0,
+        batch_size=8,
+        ddim_eta=0.0,
+        dtype=torch.float32,
+        device="cuda",
+        camera=None,
+        num_frames=4,
+        pixel_control=False,
+        transform=None,
+        offset_noise=False,
+    ):
+        """ The function supports additional image prompt.
+        Args:
+            model (_type_): the image dream model
+            image_size (_type_): size of diffusion output (standard 256)
+            prompt (_type_): text prompt for the image (prompt in type str)
+            uc (_type_): unconditional vector (tensor in shape [1, 77, 1024])
+            sampler (_type_): imagedream.ldm.models.diffusion.ddim.DDIMSampler
+            ip (Image, optional): the image prompt. Defaults to None.
+            step (int, optional): _description_. Defaults to 20.
+            scale (float, optional): _description_. Defaults to 7.5.
+            batch_size (int, optional): _description_. Defaults to 8.
+            ddim_eta (float, optional): _description_. Defaults to 0.0.
+            dtype (_type_, optional): _description_. Defaults to torch.float32.
+            device (str, optional): _description_. Defaults to "cuda".
+            camera (_type_, optional): camera info in tensor, shape: torch.Size([5, 16]) mean: 0.11, std: 0.49, min: -1.00, max: 1.00
+            num_frames (int, optional): _num of frames (views) to generate
+            pixel_control: whether to use pixel conditioning. Defaults to False, True when using pixel mode
+            transform: Compose(
+                Resize(size=(256, 256), interpolation=bilinear, max_size=None, antialias=warn)
+                ToTensor()
+                Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
+            )
+        """
+        ip_raw = ip
+        if type(prompt) != list:
+            prompt = [prompt]
+        with torch.no_grad(), torch.autocast(device_type=torch.device(device).type, dtype=dtype):
+            c = model.get_learned_conditioning(prompt).to(
+                device
+            )  # shape: torch.Size([1, 77, 1024]) mean: -0.17, std: 1.02, min: -7.50, max: 13.05
+            c_ = {"context": c.repeat(batch_size, 1, 1)}  # batch_size
+            uc_ = {"context": uc.repeat(batch_size, 1, 1)}
+
+            if camera is not None:
+                c_["camera"] = uc_["camera"] = (
+                    camera  # shape: torch.Size([5, 16]) mean: 0.11, std: 0.49, min: -1.00, max: 1.00
+                )
+                c_["num_frames"] = uc_["num_frames"] = num_frames
+
+            if ip is not None:
+                ip_embed = model.get_learned_image_conditioning(ip).to(
+                    device
+                )  # shape: torch.Size([1, 257, 1280]) mean: 0.06, std: 0.53, min: -6.83, max: 11.12
+                ip_ = ip_embed.repeat(batch_size, 1, 1)
+                c_["ip"] = ip_
+                uc_["ip"] = torch.zeros_like(ip_)
+
+            if pixel_control:
+                assert camera is not None
+                ip = transform(ip).to(
+                    device
+                )  # shape: torch.Size([3, 256, 256]) mean: 0.33, std: 0.37, min: -1.00, max: 1.00
+                ip_img = model.get_first_stage_encoding(
+                    model.encode_first_stage(ip[None, :, :, :])
+                )  # shape: torch.Size([1, 4, 32, 32]) mean: 0.23, std: 0.77, min: -4.42, max: 3.55
+                c_["ip_img"] = ip_img
+                uc_["ip_img"] = torch.zeros_like(ip_img)
+
+            shape = [4, image_size // 8, image_size // 8]  # [4, 32, 32]
+            if offset_noise:
+                ref = transform(ip_raw).to(device)
+                ref_latent = model.get_first_stage_encoding(model.encode_first_stage(ref[None, :, :, :]))
+                ref_mean = ref_latent.mean(dim=(-1, -2), keepdim=True)
+                time_steps = torch.randint(model.num_timesteps - 1, model.num_timesteps, (batch_size,), device=device)
+                x_T = model.q_sample(torch.ones([batch_size] + shape, device=device) * ref_mean, time_steps)
+
+            samples_ddim, _ = (
+                sampler.sample(  # shape: torch.Size([5, 4, 32, 32]) mean: 0.29, std: 0.85, min: -3.38, max: 4.43
+                    S=step,
+                    conditioning=c_,
+                    batch_size=batch_size,
+                    shape=shape,
+                    verbose=False,
+                    unconditional_guidance_scale=scale,
+                    unconditional_conditioning=uc_,
+                    eta=ddim_eta,
+                    x_T=x_T if offset_noise else None,
+                )
+            )
+
+            x_sample = model.decode_first_stage(samples_ddim)
+            x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
+            x_sample = 255.0 * x_sample.permute(0, 2, 3, 1).cpu().numpy()
+
+        return list(x_sample.astype(np.uint8))
+
+    def diffuse(self, t, ip, n_test=2):
+        set_seed(self.seed)
+        ip = do_resize_content(ip, self.resize_rate)
+        if self.random_background:
+            ip = add_random_background(ip)
+
+        images = []
+        for _ in range(n_test):
+            img = self.i2i(
+                self.model,
+                self.size,
+                t,
+                self.uc,
+                self.sampler,
+                ip=ip,
+                step=50,
+                scale=5,
+                batch_size=self.batch_size,
+                ddim_eta=0.0,
+                dtype=self.dtype,
+                device=self.device,
+                camera=self.camera,
+                num_frames=self.num_frames,
+                pixel_control=(self.mode == "pixel"),
+                transform=self.image_transform,
+                offset_noise=self.offset_noise,
+            )
+            img = np.concatenate(img, 1)
+            img = np.concatenate((img, ip.resize((self.size, self.size))), axis=1)
+            images.append(img)
+        set_seed()  # unset random and numpy seed
+        return images
+
+
+class ImageDreamDiffusionStage2:
+    def __init__(
+        self,
+        model,
+        device,
+        dtype,
+        num_frames,
+        camera_views,
+        ref_position,
+        random_background=False,
+        offset_noise=False,
+        resize_rate=1,
+        mode="pixel",
+        image_size=256,
+        seed=1234,
+    ) -> None:
+        assert mode in ["pixel", "local"]
+
+        size = image_size
+        self.seed = seed
+        batch_size = max(4, num_frames)
+
+        neg_texts = "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated,  obscure, unnatural colors, poor lighting, dull, and unclear."
+        uc = model.get_learned_conditioning([neg_texts]).to(device)
+        sampler = DDIMSampler(model)
+
+        # pre-compute camera matrices
+        camera = [get_camera_for_index(i).squeeze() for i in camera_views]
+        if ref_position is not None:
+            camera[ref_position] = torch.zeros_like(camera[ref_position])  # set ref camera to zero
+        camera = torch.stack(camera)
+        camera = camera.repeat(batch_size // num_frames, 1).to(device)
+
+        self.image_transform = T.Compose(
+            [
+                T.Resize((size, size)),
+                T.ToTensor(),
+                T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+            ]
+        )
+
+        self.dtype = dtype
+        self.mode = mode
+        self.ref_position = ref_position
+        self.random_background = random_background
+        self.resize_rate = resize_rate
+        self.num_frames = num_frames
+        self.size = size
+        self.device = device
+        self.batch_size = batch_size
+        self.model = model
+        self.sampler = sampler
+        self.uc = uc
+        self.camera = camera
+        self.offset_noise = offset_noise
+
+    @staticmethod
+    def i2iStage2(
+        model,
+        image_size,
+        prompt,
+        uc,
+        sampler,
+        pixel_images,
+        ip=None,
+        step=20,
+        scale=5.0,
+        batch_size=8,
+        ddim_eta=0.0,
+        dtype=torch.float32,
+        device="cuda",
+        camera=None,
+        num_frames=4,
+        pixel_control=False,
+        transform=None,
+        offset_noise=False,
+    ):
+        ip_raw = ip
+        if type(prompt) != list:
+            prompt = [prompt]
+        with torch.no_grad(), torch.autocast(device_type=torch.device(device).type, dtype=dtype):
+            c = model.get_learned_conditioning(prompt).to(
+                device
+            )  # shape: torch.Size([1, 77, 1024]) mean: -0.17, std: 1.02, min: -7.50, max: 13.05
+            c_ = {"context": c.repeat(batch_size, 1, 1)}  # batch_size
+            uc_ = {"context": uc.repeat(batch_size, 1, 1)}
+
+            if camera is not None:
+                c_["camera"] = uc_["camera"] = (
+                    camera  # shape: torch.Size([5, 16]) mean: 0.11, std: 0.49, min: -1.00, max: 1.00
+                )
+                c_["num_frames"] = uc_["num_frames"] = num_frames
+
+            if ip is not None:
+                ip_embed = model.get_learned_image_conditioning(ip).to(
+                    device
+                )  # shape: torch.Size([1, 257, 1280]) mean: 0.06, std: 0.53, min: -6.83, max: 11.12
+                ip_ = ip_embed.repeat(batch_size, 1, 1)
+                c_["ip"] = ip_
+                uc_["ip"] = torch.zeros_like(ip_)
+
+            if pixel_control:
+                assert camera is not None
+                
+            transed_pixel_images = torch.stack([transform(i).to(device) for i in pixel_images])
+            latent_pixel_images = model.get_first_stage_encoding(model.encode_first_stage(transed_pixel_images))
+
+            c_["pixel_images"] = latent_pixel_images
+            uc_["pixel_images"] = torch.zeros_like(latent_pixel_images)
+
+            shape = [4, image_size // 8, image_size // 8]  # [4, 32, 32]
+            if offset_noise:
+                ref = transform(ip_raw).to(device)
+                ref_latent = model.get_first_stage_encoding(model.encode_first_stage(ref[None, :, :, :]))
+                ref_mean = ref_latent.mean(dim=(-1, -2), keepdim=True)
+                time_steps = torch.randint(model.num_timesteps - 1, model.num_timesteps, (batch_size,), device=device)
+                x_T = model.q_sample(torch.ones([batch_size] + shape, device=device) * ref_mean, time_steps)
+
+            samples_ddim, _ = (
+                sampler.sample(  # shape: torch.Size([5, 4, 32, 32]) mean: 0.29, std: 0.85, min: -3.38, max: 4.43
+                    S=step,
+                    conditioning=c_,
+                    batch_size=batch_size,
+                    shape=shape,
+                    verbose=False,
+                    unconditional_guidance_scale=scale,
+                    unconditional_conditioning=uc_,
+                    eta=ddim_eta,
+                    x_T=x_T if offset_noise else None,
+                )
+            )
+            x_sample = model.decode_first_stage(samples_ddim)
+            x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
+            x_sample = 255.0 * x_sample.permute(0, 2, 3, 1).cpu().numpy()
+
+        return list(x_sample.astype(np.uint8))
+
+    @torch.no_grad()
+    def diffuse(self, t, ip, pixel_images, n_test=2):
+        set_seed(self.seed)
+        ip = do_resize_content(ip, self.resize_rate)
+        pixel_images = [do_resize_content(i, self.resize_rate) for i in pixel_images]
+
+        if self.random_background:
+            bg_color = np.random.rand() * 255
+            ip = add_random_background(ip, bg_color)
+            pixel_images = [add_random_background(i, bg_color) for i in pixel_images]
+
+        images = []
+        for _ in range(n_test):
+            img = self.i2iStage2(
+                self.model,
+                self.size,
+                t,
+                self.uc,
+                self.sampler,
+                pixel_images=pixel_images,
+                ip=ip,
+                step=50,
+                scale=5,
+                batch_size=self.batch_size,
+                ddim_eta=0.0,
+                dtype=self.dtype,
+                device=self.device,
+                camera=self.camera,
+                num_frames=self.num_frames,
+                pixel_control=(self.mode == "pixel"),
+                transform=self.image_transform,
+                offset_noise=self.offset_noise,
+            )
+            img = np.concatenate(img, 1)
+            img = np.concatenate(
+                (img, ip.resize((self.size, self.size)), *[i.resize((self.size, self.size)) for i in pixel_images]),
+                axis=1,
+            )
+            images.append(img)
+        set_seed()  # unset random and numpy seed
+        return images