test implementation stable cascade

handler.py

@@ -4,6 +4,7 @@ from PIL import Image
 from io import BytesIO
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from diffusers import StableDiffusionPipeline
+from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 
 import torch
 
@@ -20,6 +21,10 @@ class EndpointHandler():
          self.stable_diffusion_id = "Lykon/dreamshaper-8"
          self.pipe = StableDiffusionPipeline.from_pretrained(self.stable_diffusion_id,torch_dtype=dtype,safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to(device.type)
 
+         self.prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype)#.to(device)
+         self.decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype)#.to(device) 
+
+
          self.generator = torch.Generator(device=device.type).manual_seed(3)
 
      def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
@@ -27,25 +32,54 @@ class EndpointHandler():
 #         :param data: A dictionary contains `inputs` and optional `image` field.
 #         :return: A dictionary with `image` field contains image in base64.
 #         """
-         prompt = data.pop("inputs", None)
-         num_inference_steps = data.pop("num_inference_steps", 30)
-         guidance_scale = data.pop("guidance_scale", 7.4)
-         negative_prompt = data.pop("negative_prompt", None)
-         height = data.pop("height", None)
-         width = data.pop("width", None)
-
-        # run inference pipeline
-         out = self.pipe(
-            prompt=prompt, 
-            negative_prompt=negative_prompt,
-            num_inference_steps=num_inference_steps, 
-            guidance_scale=guidance_scale,
-            num_images_per_prompt=1,
-            height=height,
-            width=width,
-            generator=self.generator
-        )
+        prompt = data.pop("inputs", None)
+        num_inference_steps = data.pop("num_inference_steps", 30)
+        guidance_scale = data.pop("guidance_scale", 7.4)
+        negative_prompt = data.pop("negative_prompt", None)
+        height = data.pop("height", None)
+        width = data.pop("width", None)
 
+        # # run inference pipeline
+        #  out = self.pipe(
+        #     prompt=prompt, 
+        #     negative_prompt=negative_prompt,
+        #     num_inference_steps=num_inference_steps, 
+        #     guidance_scale=guidance_scale,
+        #     num_images_per_prompt=1,
+        #     height=height,
+        #     width=width,
+        #     generator=self.generator
+        # )
+        
+        self.prior_pipeline.to(device)
+        self.decoder_pipeline.to(device)
+        
+        prior_output = prior_pipeline(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        # timesteps=DEFAULT_STAGE_C_TIMESTEPS,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=1,
+        generator=self.generator,
+        # callback=callback_prior,
+        # callback_steps=callback_steps
+        )
+        
+        
+        decoder_output = self.decoder_pipeline(
+        image_embeddings=prior_output.image_embeddings,
+        prompt=prompt,
+        num_inference_steps=num_inference_steps,
+        # timesteps=decoder_timesteps,
+        guidance_scale=guidance_scale,
+        negative_prompt=negative_prompt,
+        generator=self.generator,
+        output_type="pil",
+        ).images
         
-         # return first generate PIL image
-         return out.images[0]
+        return decoder_output[0]
+
+