Merge branch 'main' of https://huggingface.co/Charles-Elena/ControlNet-endpoint-CRL-test

handler.py

@@ -4,7 +4,6 @@ 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
 
@@ -21,10 +20,6 @@ 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]]]:
@@ -32,53 +27,25 @@ 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
-        # )
-        
-        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
+         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
         )
-        
-        
-        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 decoder_output[0]
 
+        
+         # return first generate PIL image
+         return out.images[0]