add inference

app.py

@@ -198,20 +198,27 @@ def generate_image_with_steps(prompt, negative_prompt, seed, width, height, guid
                 latent_model_input  = pipe.scheduler.scale_model_input(latent_model_input , timestep=t)
                 negative_prompts = '(worst quality:2), (low quality:2), (normal quality:2), bad anatomy, bad proportions, poorly drawn face, poorly drawn hands, missing fingers, extra limbs, blurry, pixelated, distorted, lowres, jpeg artifacts, watermark, signature, text, (deformed:1.5), (bad hands:1.3), overexposed, underexposed, censored, mutated, extra fingers, cloned face, bad eyes'
                 negative_prompts = 1 * [negative_prompts]
-                        
+                use_afs = num_inference_steps < 7
+                use_free_predictor = False
                 prompt_embeds, cond_kwargs = prepare_sdxl_pipeline_step_parameter(pipe
-                                                    , prompts
-                                                    , need_cfg=True
-                                                    , device=pipe.device
-                                                    , negative_prompt=negative_prompts
-                                                    , W=width
-                                                    , H=height)
-                noise_pred  = pipe.unet(latent_model_input 
+                                                                                      , prompts
+                                                                                      , need_cfg=True
+                                                                                      , device=pipe.device
+                                                                                      , negative_prompt=negative_prompts
+                                                                                      , W=width
+                                                                                      , H=height)
+                if idx == 0 and use_afs:
+                    noise_pred = latent_model_input * 0.975
+                elif idx == len(pipe.scheduler.timesteps) - 1 and use_free_predictor:
+                    noise_pred = None
+                else:
+                    noise_pred  = pipe.unet(latent_model_input 
                                         , t
                                         , encoder_hidden_states=prompt_embeds.to(device=latents.device, dtype=latents.dtype)
                                         , added_cond_kwargs=cond_kwargs).sample
-                uncond, cond = noise_pred.chunk(2)
-                noise_pred  = uncond + (cond - uncond) * guidance_scale
+                if noise_pred is not None:
+                    uncond, cond = noise_pred.chunk(2)
+                    noise_pred  = uncond + (cond - uncond) * guidance_scale
                 latents = pipe.scheduler.step(noise_pred, t, latents).prev_sample
                 idx += 1
                         
@@ -325,13 +332,13 @@ with gr.Blocks() as demo:
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
                     minimum=0.0,
-                    maximum=10.0,
+                    maximum=6.0,
                     step=0.1,
-                    value=7.5,  # Replace with defaults that work for your model
+                    value=5.5,  # Replace with defaults that work for your model
                 )
 
                 num_inference_steps = gr.Dropdown(
-                    choices=[6, 7, 8],
+                    choices=[5, 6, 7, 8],
                     value=8,
                     label="Number of inference steps",
                 )

customed_unipc_scheduler.py

@@ -215,7 +215,8 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
         skip_type: str = "customed_time_karras",
         denoise_to_zero: bool = False,
         rescale_betas_zero_snr: bool = False,
-        use_afs: bool = False
+        use_afs: bool = False,
+        use_free_predictor = False
     ):
         
         if self.config.use_beta_sigmas and not is_scipy_available():
@@ -238,6 +239,7 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
             raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
 
         self.skip_type = skip_type
+        self.use_free_predictor = use_free_predictor
         self.use_afs = use_afs
         self.denoise_to_zero = denoise_to_zero
         if rescale_betas_zero_snr:
@@ -331,7 +333,7 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
                 ct_end = self._sigma_to_t(sigmas[9], log_sigmas)
                 if self.denoise_to_zero:
                     ct_real_end = self._sigma_to_t(sigmas[-1], log_sigmas)
-                timesteps = self.get_sigmas_karras(9, ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
+                timesteps = self.get_sigmas_karras(9 + (1 if self.use_free_predictor else 0), ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
             elif N == 5:
                 log_sigmas = np.log(sigmas)
                 sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0)
@@ -339,7 +341,7 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
                 ct_end = self._sigma_to_t(sigmas[6], log_sigmas)
                 if self.denoise_to_zero:
                     ct_real_end = self._sigma_to_t(sigmas[-1], log_sigmas)
-                timesteps = self.get_sigmas_karras(5, ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
+                timesteps = self.get_sigmas_karras(5 + (1 if self.use_afs else 0) + (1 if self.use_free_predictor else 0), ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
             elif N == 6:
                 log_sigmas = np.log(sigmas)
                 sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0)
@@ -347,7 +349,7 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
                 ct_end = self._sigma_to_t(sigmas[6], log_sigmas)
                 if self.denoise_to_zero:
                     ct_real_end = self._sigma_to_t(sigmas[-1], log_sigmas)
-                timesteps = self.get_sigmas_karras(6, ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
+                timesteps = self.get_sigmas_karras(6 + (1 if self.use_afs else 0) + (1 if self.use_free_predictor else 0), ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
             elif N == 7:
                 log_sigmas = np.log(sigmas)
                 sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0)
@@ -355,7 +357,7 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
                 ct_end = self._sigma_to_t(sigmas[6], log_sigmas)
                 if self.denoise_to_zero:
                     ct_real_end = self._sigma_to_t(sigmas[-1], log_sigmas)
-                timesteps = self.get_sigmas_karras(7, ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
+                timesteps = self.get_sigmas_karras(7 + (1 if self.use_free_predictor else 0), ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
             elif N == 8:
                 log_sigmas = np.log(sigmas).copy()
                 sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0)
@@ -363,10 +365,10 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
                 ct_end = self._sigma_to_t(sigmas[6], log_sigmas)
                 if self.denoise_to_zero:
                     ct_real_end = self._sigma_to_t(sigmas[-1], log_sigmas)
-                timesteps = self.get_sigmas_karras(8, ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
+                timesteps = self.get_sigmas_karras(8 + (1 if self.use_free_predictor else 0), ct_end, ct_start,rho=1.2, customed_final_sigma= ct_real_end if self.denoise_to_zero else None)
 
-        if self.use_afs:
-            np.insert(timesteps,1,(timesteps[0]+timesteps[1]) / 2)
+        if self.use_afs and N > 6:
+            timesteps = np.insert(timesteps,1,(timesteps[0]+timesteps[1]) / 2)
 
 
         timesteps_tmp = copy.deepcopy(timesteps)
@@ -523,6 +525,9 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
         sigma = self.sigmas[self.step_index]
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
 
+        if model_output is None:
+            return None
+
         if self.predict_x0:
             if self.config.prediction_type == "epsilon":
                 x0_pred = (sample - sigma_t * model_output) / alpha_t
@@ -560,7 +565,7 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
 
     def multistep_uni_p_bh_update(
         self,
-        model_output: torch.Tensor,
+        model_output: torch.Tensor = None,
         *args,
         sample: torch.Tensor = None,
         order: int = None,
@@ -897,7 +902,7 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
         )
 
         model_output_convert = self.convert_model_output(model_output, sample=sample)
-        if use_corrector:
+        if use_corrector and model_output_convert is not None:
             sample = self.multistep_uni_c_bh_update(
                 this_model_output=model_output_convert,
                 last_sample=self.last_sample,
@@ -908,8 +913,8 @@ class CustomedUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
         for i in range(self.solver_order - 1):
             self.model_outputs[i] = self.model_outputs[i + 1]
             self.timestep_list[i] = self.timestep_list[i + 1]
-
-        self.model_outputs[-1] = model_output_convert
+        if model_output_convert is not None:
+            self.model_outputs[-1] = model_output_convert
         self.timestep_list[-1] = timestep
 
         if self.config.lower_order_final: