Update ace_inference.py

ace_inference.py

@@ -79,153 +79,19 @@ def process_edit_image(images,
         mask_tensors.append(mask_tensor)
     return img_tensors, mask_tensors
 
-
 class TextEmbedding(nn.Module):
     def __init__(self, embedding_shape):
         super().__init__()
         self.pos = nn.Parameter(data=torch.zeros(embedding_shape))
 
-class RefinerInference(DiffusionInference):
-    def init_from_cfg(self, cfg):
-        super().init_from_cfg(cfg)
-        self.diffusion = DIFFUSIONS.build(cfg.MODEL.DIFFUSION, logger=self.logger) \
-            if cfg.MODEL.have('DIFFUSION') else None
-        self.max_seq_length = cfg.MODEL.get("MAX_SEQ_LENGTH", 4096)
-        assert self.diffusion is not None
-        self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
-        self.dynamic_load(self.diffusion_model, 'diffusion_model')
-        self.dynamic_load(self.first_stage_model, 'first_stage_model')
-
-    @torch.no_grad()
-    def encode_first_stage(self, x, **kwargs):
-        _, dtype = self.get_function_info(self.first_stage_model, 'encode')
-        with torch.autocast('cuda',
-                            enabled=dtype in ('float16', 'bfloat16'),
-                            dtype=getattr(torch, dtype)):
-            def run_one_image(u):
-                zu = get_model(self.first_stage_model).encode(u)
-                if isinstance(zu, (tuple, list)):
-                    zu = zu[0]
-                return zu
-            z = [run_one_image(u.unsqueeze(0) if u.dim == 3 else u) for u in x]
-            return z
-    def upscale_resize(self, image, interpolation=T.InterpolationMode.BILINEAR):
-        c, H, W = image.shape
-        scale = max(1.0, math.sqrt(self.max_seq_length / ((H / 16) * (W / 16))))
-        rH = int(H * scale) // 16 * 16  # ensure divisible by self.d
-        rW = int(W * scale) // 16 * 16
-        image = T.Resize((rH, rW), interpolation=interpolation, antialias=True)(image)
-        return image
-    @torch.no_grad()
-    def decode_first_stage(self, z):
-        _, dtype = self.get_function_info(self.first_stage_model, 'decode')
-        with torch.autocast('cuda',
-                            enabled=dtype in ('float16', 'bfloat16'),
-                            dtype=getattr(torch, dtype)):
-            return [get_model(self.first_stage_model).decode(zu) for zu in z]
-
-    def noise_sample(self, num_samples, h, w, seed, device = None, dtype = torch.bfloat16):
-        noise = torch.randn(
-            num_samples,
-            16,
-            # allow for packing
-            2 * math.ceil(h / 16),
-            2 * math.ceil(w / 16),
-            device=device,
-            dtype=dtype,
-            generator=torch.Generator(device=device).manual_seed(seed),
-        )
-        return noise
-    def refine(self,
-               x_samples=None,
-               prompt=None,
-               reverse_scale=-1.,
-               seed = 2024,
-               use_dynamic_model = False,
-               **kwargs
-               ):
-        print(prompt)
-        value_input = copy.deepcopy(self.input)
-        x_samples = [self.upscale_resize(x) for x in x_samples]
-
-        noise = []
-        for i, x in enumerate(x_samples):
-            noise_ = self.noise_sample(1, x.shape[1],
-                                       x.shape[2], seed,
-                                       device = x.device)
-            noise.append(noise_)
-        noise, x_shapes = pack_imagelist_into_tensor(noise)
-        if reverse_scale > 0:
-            if use_dynamic_model: self.dynamic_load(self.first_stage_model, 'first_stage_model')
-            x_samples = [x.unsqueeze(0) for x in x_samples]
-            x_start = self.encode_first_stage(x_samples, **kwargs)
-            if use_dynamic_model: self.dynamic_unload(self.first_stage_model,
-                                'first_stage_model',
-                                skip_loaded=True)
-            x_start, _ = pack_imagelist_into_tensor(x_start)
-        else:
-            x_start = None
-        # cond stage
-        if use_dynamic_model: self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
-        function_name, dtype = self.get_function_info(self.cond_stage_model)
-        with torch.autocast('cuda',
-                            enabled=dtype == 'float16',
-                            dtype=getattr(torch, dtype)):
-            ctx = getattr(get_model(self.cond_stage_model),
-                          function_name)(prompt)
-            ctx["x_shapes"] = x_shapes
-        if use_dynamic_model: self.dynamic_unload(self.cond_stage_model,
-                            'cond_stage_model',
-                            skip_loaded=True)
-
-
-        if use_dynamic_model: self.dynamic_load(self.diffusion_model, 'diffusion_model')
-        # UNet use input n_prompt
-        function_name, dtype = self.get_function_info(
-            self.diffusion_model)
-        with torch.autocast('cuda',
-                            enabled=dtype in ('float16', 'bfloat16'),
-                            dtype=getattr(torch, dtype)):
-            solver_sample = value_input.get('sample', 'flow_euler')
-            sample_steps = value_input.get('sample_steps', 20)
-            guide_scale = value_input.get('guide_scale', 3.5)
-            if guide_scale is not None:
-                guide_scale = torch.full((noise.shape[0],), guide_scale, device=noise.device,
-                                         dtype=noise.dtype)
-            else:
-                guide_scale = None
-            latent = self.diffusion.sample(
-                noise=noise,
-                sampler=solver_sample,
-                model=get_model(self.diffusion_model),
-                model_kwargs={"cond": ctx, "guidance": guide_scale},
-                steps=sample_steps,
-                show_progress=True,
-                guide_scale=guide_scale,
-                return_intermediate=None,
-                reverse_scale=reverse_scale,
-                x=x_start,
-                **kwargs).float()
-        latent = unpack_tensor_into_imagelist(latent, x_shapes)
-        if use_dynamic_model: self.dynamic_unload(self.diffusion_model,
-                            'diffusion_model',
-                            skip_loaded=True)
-        if use_dynamic_model: self.dynamic_load(self.first_stage_model, 'first_stage_model')
-        x_samples = self.decode_first_stage(latent)
-        if use_dynamic_model: self.dynamic_unload(self.first_stage_model,
-                            'first_stage_model',
-                            skip_loaded=True)
-        return x_samples
-
-
-class ACEInference(DiffusionInference):
+class ACEFluxLCInference(DiffusionInference):
     def __init__(self, logger=None):
         if logger is None:
             logger = get_logger(name='scepter')
         self.logger = logger
         self.loaded_model = {}
         self.loaded_model_name = [
-            'diffusion_model', 'first_stage_model', 'cond_stage_model'
+            'diffusion_model', 'first_stage_model', 'cond_stage_model', 'ref_cond_stage_model'
         ]
 
     def init_from_cfg(self, cfg):
@@ -234,7 +100,7 @@ class ACEInference(DiffusionInference):
         self.use_dynamic_model = cfg.get('USE_DYNAMIC_MODEL', True)
         module_paras = self.load_default(cfg.get('DEFAULT_PARAS', None))
         assert cfg.have('MODEL')
-
+        self.size_factor = cfg.get('SIZE_FACTOR', 8)
         self.diffusion_model = self.infer_model(
             cfg.MODEL.DIFFUSION_MODEL, module_paras.get(
                 'DIFFUSION_MODEL',
@@ -250,24 +116,23 @@ class ACEInference(DiffusionInference):
                 'COND_STAGE_MODEL',
                 None)) if cfg.MODEL.have('COND_STAGE_MODEL') else None
 
-        self.refiner_model_cfg = cfg.get('REFINER_MODEL', None)
-        # self.refiner_scale = cfg.get('REFINER_SCALE', 0.)
-        # self.refiner_prompt = cfg.get('REFINER_PROMPT', "")
-        self.ace_prompt = cfg.get("ACE_PROMPT", [])
-        if self.refiner_model_cfg:
-            self.refiner_module = RefinerInference(self.logger)
-            self.refiner_module.init_from_cfg(self.refiner_model_cfg)
-        else:
-            self.refiner_module = None
+        self.ref_cond_stage_model = self.infer_model(
+            cfg.MODEL.REF_COND_STAGE_MODEL,
+            module_paras.get(
+                'REF_COND_STAGE_MODEL',
+                None)) if cfg.MODEL.have('REF_COND_STAGE_MODEL') else None
 
         self.diffusion = DIFFUSIONS.build(cfg.MODEL.DIFFUSION,
                                           logger=self.logger)
-
-
         self.interpolate_func = lambda x: (F.interpolate(
             x.unsqueeze(0),
             scale_factor=1 / self.size_factor,
             mode='nearest-exact') if x is not None else None)
+
+        self.max_seq_length = cfg.get("MAX_SEQ_LENGTH", 4096)
+        self.src_max_seq_length = cfg.get("SRC_MAX_SEQ_LENGTH", 1024)
+        self.image_token = cfg.MODEL.get("IMAGE_TOKEN", "<img>")
+
         self.text_indentifers = cfg.MODEL.get('TEXT_IDENTIFIER', [])
         self.use_text_pos_embeddings = cfg.MODEL.get('USE_TEXT_POS_EMBEDDINGS',
                                                      False)
@@ -277,41 +142,66 @@ class ACEInference(DiffusionInference):
         else:
             self.text_position_embeddings = None
 
-        self.max_seq_len = cfg.MODEL.DIFFUSION_MODEL.MAX_SEQ_LEN
-        self.scale_factor = cfg.get('SCALE_FACTOR', 0.18215)
-        self.size_factor = cfg.get('SIZE_FACTOR', 8)
-        self.decoder_bias = cfg.get('DECODER_BIAS', 0)
-        self.default_n_prompt = cfg.get('DEFAULT_N_PROMPT', '')
-        #self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
-        #self.dynamic_load(self.diffusion_model, 'diffusion_model')
-        #self.dynamic_load(self.first_stage_model, 'first_stage_model')
+        if not self.use_dynamic_model:
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+            if self.ref_cond_stage_model is not None: self.dynamic_load(self.ref_cond_stage_model, 'ref_cond_stage_model')
+            self.dynamic_load(self.diffusion_model, 'diffusion_model')
+
+    def upscale_resize(self, image, interpolation=T.InterpolationMode.BILINEAR):
+        c, H, W = image.shape
+        scale = max(1.0, math.sqrt(self.max_seq_length / ((H / 16) * (W / 16))))
+        rH = int(H * scale) // 16 * 16  # ensure divisible by self.d
+        rW = int(W * scale) // 16 * 16
+        image = T.Resize((rH, rW), interpolation=interpolation, antialias=True)(image)
+        return image
+
 
     @torch.no_grad()
     def encode_first_stage(self, x, **kwargs):
         _, dtype = self.get_function_info(self.first_stage_model, 'encode')
         with torch.autocast('cuda',
-                            enabled=(dtype != 'float32'),
+                            enabled=dtype in ('float16', 'bfloat16'),
                             dtype=getattr(torch, dtype)):
-            z = [
-                self.scale_factor * get_model(self.first_stage_model)._encode(
-                    i.unsqueeze(0).to(getattr(torch, dtype))) for i in x
-            ]
-        return z
+            def run_one_image(u):
+                zu = get_model(self.first_stage_model).encode(u)
+                if isinstance(zu, (tuple, list)):
+                    zu = zu[0]
+                return zu
+
+            z = [run_one_image(u.unsqueeze(0) if u.dim() == 3 else u) for u in x]
+            return z
+
 
     @torch.no_grad()
     def decode_first_stage(self, z):
         _, dtype = self.get_function_info(self.first_stage_model, 'decode')
         with torch.autocast('cuda',
-                            enabled=(dtype != 'float32'),
+                            enabled=dtype in ('float16', 'bfloat16'),
                             dtype=getattr(torch, dtype)):
-            x = [
-                get_model(self.first_stage_model)._decode(
-                    1. / self.scale_factor * i.to(getattr(torch, dtype)))
-                for i in z
-            ]
-        return x
+            return [get_model(self.first_stage_model).decode(zu) for zu in z]
 
+    def noise_sample(self, num_samples, h, w, seed, device = None, dtype = torch.bfloat16):
+        noise = torch.randn(
+            num_samples,
+            16,
+            # allow for packing
+            2 * math.ceil(h / 16),
+            2 * math.ceil(w / 16),
+            device=device,
+            dtype=dtype,
+            generator=torch.Generator(device=device).manual_seed(seed),
+        )
+        return noise
 
+    # def preprocess_prompt(self, prompt):
+    #     prompt_ = [[pp] if isinstance(pp, str) else pp for pp in prompt]
+    #     for pp_id, pp in enumerate(prompt_):
+    #         prompt_[pp_id] = [""] + pp
+    #         for p_id, p in enumerate(prompt_[pp_id]):
+    #             prompt_[pp_id][p_id] = self.image_token + self.text_indentifers[p_id] + " " + p
+    #         prompt_[pp_id] = [f";".join(prompt_[pp_id])]
+    #     return prompt_
 
     @torch.no_grad()
     def __call__(self,
@@ -320,48 +210,35 @@ class ACEInference(DiffusionInference):
                  prompt='',
                  task=None,
                  negative_prompt='',
-                 output_height=512,
-                 output_width=512,
-                 sampler='ddim',
+                 output_height=1024,
+                 output_width=1024,
+                 sampler='flow_euler',
                  sample_steps=20,
-                 guide_scale=4.5,
-                 guide_rescale=0.5,
+                 guide_scale=3.5,
                  seed=-1,
                  history_io=None,
                  tar_index=0,
+                 align=0,
                  **kwargs):
         input_image, input_mask = image, mask
-        g = torch.Generator(device=we.device_id)
         seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
-        g.manual_seed(int(seed))
         if input_image is not None:
             # assert isinstance(input_image, list) and isinstance(input_mask, list)
             if task is None:
                 task = [''] * len(input_image)
             if not isinstance(prompt, list):
                 prompt = [prompt] * len(input_image)
-            if history_io is not None and len(history_io) > 0:
-                his_image, his_maks, his_prompt, his_task = history_io[
-                    'image'], history_io['mask'], history_io[
-                        'prompt'], history_io['task']
-                assert len(his_image) == len(his_maks) == len(
-                    his_prompt) == len(his_task)
-                input_image = his_image + input_image
-                input_mask = his_maks + input_mask
-                task = his_task + task
-                prompt = his_prompt + [prompt[-1]]
-                prompt = [
-                    pp.replace('{image}', f'{{image{i}}}') if i > 0 else pp
-                    for i, pp in enumerate(prompt)
-                ]
-
+            prompt = [
+                pp.replace('{image}', f'{{image{i}}}') if i > 0 else pp
+                for i, pp in enumerate(prompt)
+            ]
             edit_image, edit_image_mask = process_edit_image(
-                input_image, input_mask, task, max_seq_len=self.max_seq_len)
-
-            image, image_mask = edit_image[tar_index], edit_image_mask[
-                tar_index]
+                input_image, input_mask, task, max_seq_len=self.src_max_seq_length)
+            image, image_mask = self.upscale_resize(edit_image[tar_index]), self.upscale_resize(edit_image_mask[
+               tar_index])
+            # edit_image, edit_image_mask = [[self.upscale_resize(i) for i in edit_image]], [[self.upscale_resize(i) for i in edit_image_mask]]
+            # image, image_mask = edit_image[tar_index], edit_image_mask[tar_index]
             edit_image, edit_image_mask = [edit_image], [edit_image_mask]
-
         else:
             edit_image = edit_image_mask = [[]]
             image = torch.zeros(
@@ -373,177 +250,107 @@ class ACEInference(DiffusionInference):
             if not isinstance(prompt, list):
                 prompt = [prompt]
 
-        image, image_mask, prompt = [image], [image_mask], [prompt]
+        image, image_mask, prompt = [image], [image_mask], [prompt],
+        align = [align for p in prompt] if isinstance(align, int) else align
+
         assert check_list_of_list(prompt) and check_list_of_list(
             edit_image) and check_list_of_list(edit_image_mask)
-        # Assign Negative Prompt
-        if isinstance(negative_prompt, list):
-            negative_prompt = negative_prompt[0]
-        assert isinstance(negative_prompt, str)
-
-        n_prompt = copy.deepcopy(prompt)
-        for nn_p_id, nn_p in enumerate(n_prompt):
-            assert isinstance(nn_p, list)
-            n_prompt[nn_p_id][-1] = negative_prompt
-
-        is_txt_image = sum([len(e_i) for e_i in edit_image]) < 1
+        # negative prompt is not used
         image = to_device(image)
+        ctx = {}
+        # Get Noise Shape
+        self.dynamic_load(self.first_stage_model, 'first_stage_model')
+        x = self.encode_first_stage(image)
+        self.dynamic_unload(self.first_stage_model,
+                            'first_stage_model',
+                            skip_loaded=not self.use_dynamic_model)
 
-        refiner_scale = kwargs.pop("refiner_scale", 0.0)
-        refiner_prompt = kwargs.pop("refiner_prompt", "")
-        use_ace = kwargs.pop("use_ace", True)
-        # <= 0 use ace as the txt2img generator.
-        if use_ace and (not is_txt_image or refiner_scale <= 0):
-            ctx, null_ctx = {}, {}
-            # Get Noise Shape
-            self.dynamic_load(self.first_stage_model, 'first_stage_model')
-            x = self.encode_first_stage(image)
-            self.dynamic_unload(self.first_stage_model,
-                                'first_stage_model',
-                                skip_loaded=True)
-            noise = [
-                torch.empty(*i.shape, device=we.device_id).normal_(generator=g)
-                for i in x
-            ]
-            noise, x_shapes = pack_imagelist_into_tensor(noise)
-            ctx['x_shapes'] = null_ctx['x_shapes'] = x_shapes
+        g = torch.Generator(device=we.device_id).manual_seed(seed)
 
-            image_mask = to_device(image_mask, strict=False)
-            cond_mask = [self.interpolate_func(i) for i in image_mask
-                         ] if image_mask is not None else [None] * len(image)
-            ctx['x_mask'] = null_ctx['x_mask'] = cond_mask
+        noise = [
+            torch.randn((1, 16, i.shape[2], i.shape[3]), device=we.device_id, dtype=torch.bfloat16).normal_(generator=g)
+            for i in x
+        ]
+        noise, x_shapes = pack_imagelist_into_tensor(noise)
+        ctx['x_shapes'] = x_shapes
+        ctx['align'] = align
+
+        image_mask = to_device(image_mask, strict=False)
+        cond_mask = [self.interpolate_func(i) for i in image_mask
+                     ] if image_mask is not None else [None] * len(image)
+        ctx['x_mask'] = cond_mask
+        # Encode Prompt
+        instruction_prompt = [[pp[-1]] if "{image}" in pp[-1] else ["{image} " + pp[-1]] for pp in prompt]
+        self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+        function_name, dtype = self.get_function_info(self.cond_stage_model)
+        cont = getattr(get_model(self.cond_stage_model), function_name)(instruction_prompt)
+        cont["context"] = [ct[-1] for ct in cont["context"]]
+        cont["y"] = [ct[-1] for ct in cont["y"]]
+        self.dynamic_unload(self.cond_stage_model,
+                            'cond_stage_model',
+                            skip_loaded=not self.use_dynamic_model)
+        ctx.update(cont)
 
-            # Encode Prompt
-            self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
-            function_name, dtype = self.get_function_info(self.cond_stage_model)
-            cont, cont_mask = getattr(get_model(self.cond_stage_model),
-                                      function_name)(prompt)
-            cont, cont_mask = self.cond_stage_embeddings(prompt, edit_image, cont,
-                                                         cont_mask)
-            null_cont, null_cont_mask = getattr(get_model(self.cond_stage_model),
-                                                function_name)(n_prompt)
-            null_cont, null_cont_mask = self.cond_stage_embeddings(
-                prompt, edit_image, null_cont, null_cont_mask)
-            self.dynamic_unload(self.cond_stage_model,
-                                'cond_stage_model',
-                                skip_loaded=False)
-            ctx['crossattn'] = cont
-            null_ctx['crossattn'] = null_cont
-
-            # Encode Edit Images
-            self.dynamic_load(self.first_stage_model, 'first_stage_model')
-            edit_image = [to_device(i, strict=False) for i in edit_image]
-            edit_image_mask = [to_device(i, strict=False) for i in edit_image_mask]
-            e_img, e_mask = [], []
-            for u, m in zip(edit_image, edit_image_mask):
-                if u is None:
-                    continue
-                if m is None:
-                    m = [None] * len(u)
-                e_img.append(self.encode_first_stage(u, **kwargs))
-                e_mask.append([self.interpolate_func(i) for i in m])
-            self.dynamic_unload(self.first_stage_model,
-                                'first_stage_model',
-                                skip_loaded=True)
-            null_ctx['edit'] = ctx['edit'] = e_img
-            null_ctx['edit_mask'] = ctx['edit_mask'] = e_mask
-
-            # Diffusion Process
-            self.dynamic_load(self.diffusion_model, 'diffusion_model')
-            function_name, dtype = self.get_function_info(self.diffusion_model)
-            with torch.autocast('cuda',
-                                enabled=dtype in ('float16', 'bfloat16'),
-                                dtype=getattr(torch, dtype)):
-                latent = self.diffusion.sample(
-                    noise=noise,
-                    sampler=sampler,
-                    model=get_model(self.diffusion_model),
-                    model_kwargs=[{
-                        'cond':
-                        ctx,
-                        'mask':
-                        cont_mask,
-                        'text_position_embeddings':
-                        self.text_position_embeddings.pos if hasattr(
-                            self.text_position_embeddings, 'pos') else None
-                    }, {
-                        'cond':
-                        null_ctx,
-                        'mask':
-                        null_cont_mask,
-                        'text_position_embeddings':
-                        self.text_position_embeddings.pos if hasattr(
-                            self.text_position_embeddings, 'pos') else None
-                    }] if guide_scale is not None and guide_scale > 1 else {
-                        'cond':
-                        null_ctx,
-                        'mask':
-                        cont_mask,
-                        'text_position_embeddings':
-                        self.text_position_embeddings.pos if hasattr(
-                            self.text_position_embeddings, 'pos') else None
-                    },
-                    steps=sample_steps,
-                    show_progress=True,
-                    seed=seed,
-                    guide_scale=guide_scale,
-                    guide_rescale=guide_rescale,
-                    return_intermediate=None,
-                    **kwargs)
-            self.dynamic_unload(self.diffusion_model,
-                                'diffusion_model',
-                                skip_loaded=False)
-
-            # Decode to Pixel Space
-            self.dynamic_load(self.first_stage_model, 'first_stage_model')
-            samples = unpack_tensor_into_imagelist(latent, x_shapes)
-            x_samples = self.decode_first_stage(samples)
-            self.dynamic_unload(self.first_stage_model,
-                                'first_stage_model',
-                                skip_loaded=False)
-            x_samples = [x.squeeze(0) for x in x_samples]
+        # Encode Edit Images
+        self.dynamic_load(self.first_stage_model, 'first_stage_model')
+        edit_image = [to_device(i, strict=False) for i in edit_image]
+        edit_image_mask = [to_device(i, strict=False) for i in edit_image_mask]
+        e_img, e_mask = [], []
+        for u, m in zip(edit_image, edit_image_mask):
+            if u is None:
+                continue
+            if m is None:
+                m = [None] * len(u)
+            e_img.append(self.encode_first_stage(u, **kwargs))
+            e_mask.append([self.interpolate_func(i) for i in m])
+        self.dynamic_unload(self.first_stage_model,
+                            'first_stage_model',
+                            skip_loaded=not self.use_dynamic_model)
+        ctx['edit_x'] = e_img
+        ctx['edit_mask'] = e_mask
+        # Encode Ref Images
+        if guide_scale is not None:
+            guide_scale = torch.full((noise.shape[0],), guide_scale, device=noise.device, dtype=noise.dtype)
         else:
-            x_samples = image
-        if self.refiner_module and refiner_scale > 0:
-            if is_txt_image:
-                random.shuffle(self.ace_prompt)
-                input_refine_prompt = [self.ace_prompt[0] + refiner_prompt if p[0] == "" else p[0] for p in prompt]
-                input_refine_scale = -1.
-            else:
-                input_refine_prompt = [p[0].replace("{image}", "") + " " + refiner_prompt for p in prompt]
-                input_refine_scale = refiner_scale
-                print(input_refine_prompt)
-
-            x_samples = self.refiner_module.refine(x_samples,
-                                                   reverse_scale = input_refine_scale,
-                                                   prompt= input_refine_prompt,
-                                                   seed=seed,
-                                                   use_dynamic_model=self.use_dynamic_model)
+            guide_scale = None
+
+        # Diffusion Process
+        self.dynamic_load(self.diffusion_model, 'diffusion_model')
+        function_name, dtype = self.get_function_info(self.diffusion_model)
+        with torch.autocast('cuda',
+                            enabled=dtype in ('float16', 'bfloat16'),
+                            dtype=getattr(torch, dtype)):
+            latent = self.diffusion.sample(
+                noise=noise,
+                sampler=sampler,
+                model=get_model(self.diffusion_model),
+                model_kwargs={
+                    "cond": ctx, "guidance": guide_scale, "gc_seg": -1
+                },
+                steps=sample_steps,
+                show_progress=True,
+                guide_scale=guide_scale,
+                return_intermediate=None,
+                reverse_scale=-1,
+                **kwargs).float()
+        if self.use_dynamic_model: self.dynamic_unload(self.diffusion_model,
+                            'diffusion_model',
+                            skip_loaded=not self.use_dynamic_model)
+
+        # Decode to Pixel Space
+        self.dynamic_load(self.first_stage_model, 'first_stage_model')
+        samples = unpack_tensor_into_imagelist(latent, x_shapes)
+        x_samples = self.decode_first_stage(samples)
+        self.dynamic_unload(self.first_stage_model,
+                            'first_stage_model',
+                            skip_loaded=not self.use_dynamic_model)
+        x_samples = [x.squeeze(0) for x in x_samples]
 
         imgs = [
-            torch.clamp((x_i.float() + 1.0) / 2.0 + self.decoder_bias / 255,
+            torch.clamp((x_i.float() + 1.0) / 2.0,
                         min=0.0,
                         max=1.0).squeeze(0).permute(1, 2, 0).cpu().numpy()
             for x_i in x_samples
         ]
         imgs = [Image.fromarray((img * 255).astype(np.uint8)) for img in imgs]
         return imgs
-
-    def cond_stage_embeddings(self, prompt, edit_image, cont, cont_mask):
-        if self.use_text_pos_embeddings and not torch.sum(
-                self.text_position_embeddings.pos) > 0:
-            identifier_cont, _ = getattr(get_model(self.cond_stage_model),
-                                         'encode')(self.text_indentifers,
-                                                   return_mask=True)
-            self.text_position_embeddings.load_state_dict(
-                {'pos': identifier_cont[:, 0, :]})
-
-        cont_, cont_mask_ = [], []
-        for pp, edit, c, cm in zip(prompt, edit_image, cont, cont_mask):
-            if isinstance(pp, list):
-                cont_.append([c[-1], *c] if len(edit) > 0 else [c[-1]])
-                cont_mask_.append([cm[-1], *cm] if len(edit) > 0 else [cm[-1]])
-            else:
-                raise NotImplementedError
-
-        return cont_, cont_mask_