test refeactor

pipeline.py

@@ -673,6 +673,7 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         eta: float = 0.0,
         generator: Optional[torch.Generator] = None,
         latents: Optional[torch.FloatTensor] = None,
+        return_latents: bool = False,
         max_embeddings_multiples: Optional[int] = 3,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
@@ -796,7 +797,7 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device, image is None)
         latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
 
-        # 6. Prepare latent variables
+
         latents, init_latents_orig, noise = self.prepare_latents(
             image,
             latent_timestep,
@@ -812,14 +813,14 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
         # print("before denoise latents")
-        print(latents.shape)
+        # print(latents.shape)
         # 8. Denoising loop
         for i, t in enumerate(self.progress_bar(timesteps)):
             # expand the latents if we are doing classifier free guidance
             latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-            print("latent_model_input 1st step")
+            # print("latent_model_input 1st step")
             # print(latent_model_input)
-            print(latent_model_input.shape)
+            # print(latent_model_input.shape)
 
             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
 
@@ -845,86 +846,91 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
                     callback(i, t, latents)
                 if is_cancelled_callback is not None and is_cancelled_callback():
                     return None
-        print("after first step denoise latents")
+        # print("after first step denoise latents")
         # print(latents)
-        print(latents.shape)
-        upscale_latents = torch.nn.functional.interpolate(
-            latents, size=(int(height*resize_scale)//8, int(width*resize_scale)//8))
+        # print(latents.shape)
+        # upscale_latents = torch.nn.functional.interpolate(
+        #     latents, size=(int(height*resize_scale)//8, int(width*resize_scale)//8))
         
-        for i, t in enumerate(self.progress_bar(timesteps)):
-            # expand the latents if we are doing classifier free guidance
-            latent_model_input = torch.cat(
-                [upscale_latents] * 2) if do_classifier_free_guidance else upscale_latents
-            print("latent_model_input 2nd step")
-            # print(latent_model_input)
-            print(latent_model_input.shape)
-
-            print("2nd step timestep")
-            print(t)
-
-            latent_model_input = self.scheduler.scale_model_input(
-                latent_model_input, t)
+        # for i, t in enumerate(self.progress_bar(timesteps)):
+        #     # expand the latents if we are doing classifier free guidance
+        #     latent_model_input = torch.cat(
+        #         [upscale_latents] * 2) if do_classifier_free_guidance else upscale_latents
+        #     # print("latent_model_input 2nd step")
+        #     # print(latent_model_input)
+        #     # print(latent_model_input.shape)
+
+        #     # print("2nd step timestep")
+        #     # print(t)
+
+        #     latent_model_input = self.scheduler.scale_model_input(
+        #         latent_model_input, t)
             
 
 
 
-            print("latent_model_input after scheduler")
-            # print(latent_model_input)
-            print(latent_model_input.shape)
-            # predict the noise residual
-            upscale_noise_pred = self.unet(latent_model_input, t,
-                                   encoder_hidden_states=text_embeddings).sample
-            print("noise_pred")
-            # print(noise_pred)
-            print(upscale_noise_pred.shape)
-
-            print("perform guidance")
-            # perform guidance
-            if do_classifier_free_guidance:
-                noise_pred_uncond, noise_pred_text = upscale_noise_pred.chunk(
-                    2)
-                upscale_noise_pred = noise_pred_uncond + guidance_scale * \
-                    (noise_pred_text - noise_pred_uncond)
-            print("noise_pred after guidance")
-            print(upscale_noise_pred.shape)
-
-            print("compute the previous noisy sample")
-            # compute the previous noisy sample x_t -> x_t-1
-            upscale_latents = self.scheduler.step(
-                upscale_noise_pred, t, upscale_latents, **extra_step_kwargs).prev_sample
-            print(upscale_latents.shape)
-
-            print("compute mask")
-            if mask is not None:
-                # masking
-                init_latents_proper = self.scheduler.add_noise(
-                    init_latents_orig, noise, torch.tensor([t]))
-                upscale_latents = (init_latents_proper *
-                                   mask) + (latents * (1 - mask))
-
-            # call the callback, if provided
-            if i % callback_steps == 0:
-                if callback is not None:
-                    callback(i, t, upscale_latents)
-                if is_cancelled_callback is not None and is_cancelled_callback():
-                    return None
-        #do latent upscale  here
-
-        # 9. Post-processing
-        image = self.decode_latents(upscale_latents)
+        #     # print("latent_model_input after scheduler")
+        #     # print(latent_model_input)
+        #     # print(latent_model_input.shape)
+        #     # predict the noise residual
+        #     upscale_noise_pred = self.unet(latent_model_input, t,
+        #                            encoder_hidden_states=text_embeddings).sample
+        #     # print("noise_pred")
+        #     # print(noise_pred)
+        #     # print(upscale_noise_pred.shape)
+
+        #     # print("perform guidance")
+        #     # perform guidance
+        #     if do_classifier_free_guidance:
+        #         noise_pred_uncond, noise_pred_text = upscale_noise_pred.chunk(
+        #             2)
+        #         upscale_noise_pred = noise_pred_uncond + guidance_scale * \
+        #             (noise_pred_text - noise_pred_uncond)
+        #     # print("noise_pred after guidance")
+        #     # print(upscale_noise_pred.shape)
+
+        #     # print("compute the previous noisy sample")
+        #     # compute the previous noisy sample x_t -> x_t-1
+        #     upscale_latents = self.scheduler.step(
+        #         noise_pred, t, upscale_latents, **extra_step_kwargs).prev_sample
+        #     # print(upscale_latents.shape)
+
+        #     # print("compute mask")
+        #     if mask is not None:
+        #         # masking
+        #         init_latents_proper = self.scheduler.add_noise(
+        #             init_latents_orig, noise, torch.tensor([t]))
+        #         upscale_latents = (init_latents_proper *
+        #                            mask) + (latents * (1 - mask))
+
+        #     # call the callback, if provided
+        #     if i % callback_steps == 0:
+        #         if callback is not None:
+        #             callback(i, t, upscale_latents)
+        #         if is_cancelled_callback is not None and is_cancelled_callback():
+        #             return None
+        # #do latent upscale  here
+
+        # # 9. Post-processing
+        # image = self.decode_latents(upscale_latents)
+        
+        if return_latents:
+            return latents
+        else:
+            image = self.decode_latents(latents)
 
 
-        # 10. Run safety checker
-        image, has_nsfw_concept = self.run_safety_checker(image, device, text_embeddings.dtype)
+            # 10. Run safety checker
+            image, has_nsfw_concept = self.run_safety_checker(image, device, text_embeddings.dtype)
 
-        # 11. Convert to PIL
-        if output_type == "pil":
-            image = self.numpy_to_pil(image)
+            # 11. Convert to PIL
+            if output_type == "pil":
+                image = self.numpy_to_pil(image)
 
-        if not return_dict:
-            return image, has_nsfw_concept
+            if not return_dict:
+                return image, has_nsfw_concept
 
-        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+            return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
 
     def text2img(
         self,
@@ -934,6 +940,8 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         width: int = 512,
         num_inference_steps: int = 50,
         guidance_scale: float = 7.5,
+        strength: float = 0.6,
+        resize_scale: float = 1.2,
         num_images_per_prompt: Optional[int] = 1,
         eta: float = 0.0,
         generator: Optional[torch.Generator] = None,
@@ -1002,7 +1010,8 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
             (nsfw) content, according to the `safety_checker`.
         """
-        return self.__call__(
+
+        latents = self.__call__(
             prompt=prompt,
             negative_prompt=negative_prompt,
             height=height,
@@ -1013,6 +1022,30 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             eta=eta,
             generator=generator,
             latents=latents,
+            return_latents=True,
+            max_embeddings_multiples=max_embeddings_multiples,
+            output_type=output_type,
+            return_dict=return_dict,
+            callback=callback,
+            is_cancelled_callback=is_cancelled_callback,
+            callback_steps=callback_steps,
+        )
+
+        latents = torch.nn.functional.interpolate(latents, size=(int(height*resize_scale)//8, int(width*resize_scale)//8))
+
+        return self.__call__(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height*resize_scale,
+            width=width*resize_scale,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            strength=strength,
+            num_images_per_prompt=num_images_per_prompt,
+            eta=eta,
+            generator=generator,
+            latents=latents,
+            return_latents=False,
             max_embeddings_multiples=max_embeddings_multiples,
             output_type=output_type,
             return_dict=return_dict,