test basic stable cascade

handler.py

@@ -3,10 +3,9 @@ import base64
 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
+from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 
 
 # # set device
@@ -17,56 +16,45 @@ if device.type != 'cuda':
 dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
 
 class EndpointHandler():
-     def __init__(self, path=""):
-         self.stable_diffusion_id = "Lykon/dreamshaper-8"
+    #  def __init__(self, path=""):
+        #  self.stable_diffusion_id = "Lykon/dreamshaper-8"
 
-         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.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)
+        #  self.generator = torch.Generator(device=device.type).manual_seed(3)
 
      def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
-#         """
-#         :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)
+#      import torch
+
+        device = "cuda"
+        num_images_per_prompt = 2
+
+        prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16).to(device)
+        decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=torch.float16).to(device)
+
+        prompt = "Anthropomorphic cat dressed as a pilot"
+        negative_prompt = ""
 
-        self.prior_pipeline.to(device)
-        self.decoder_pipeline.to(device)
-        
-        prior_output = self.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
+        prior_output = prior(
+            prompt=prompt,
+            height=1024,
+            width=1024,
+            negative_prompt=negative_prompt,
+            guidance_scale=4.0,
+            num_images_per_prompt=num_images_per_prompt,
+            num_inference_steps=20
         )
-        
-        
-        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",
+        decoder_output = decoder(
+            image_embeddings=prior_output.image_embeddings.half(),
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=0.0,
+            output_type="pil",
+            num_inference_steps=10
         ).images
-        
-        return decoder_output[0]
+        return decoder_output
+