Update handler.py

handler.py

@@ -11,26 +11,26 @@ class EndpointHandler:
         self.model_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
 
         # # Load the VAE model
-        # self.vae = AutoencoderKL.from_pretrained(self.v_autoencoder, torch_dtype=torch.float16)
+        self.vae = AutoencoderKL.from_pretrained(self.v_autoencoder, torch_dtype=torch.float16)
 
-        # # Load the main pipeline
-        # self.pipe = StableDiffusionXLPipeline.from_pretrained(
-        #     self.model_base,
-        #     torch_dtype=torch.float16,
-        #     vae=self.vae,
-        #     add_watermarker=False,
-        # )
-        # self.pipe.safety_checker = None
-        # self.pipe.to("cuda")
+        # Load the main pipeline
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            self.model_base,
+            torch_dtype=torch.float16,
+            vae=self.vae,
+            add_watermarker=False,
+        )
+        self.pipe.safety_checker = None
+        self.pipe.to("cuda")
 
-        # # Load the refiner pipeline
-        # self.pipe_refiner = DiffusionPipeline.from_pretrained(
-        #     self.model_refiner, 
-        #     torch_dtype=torch.float16, 
-        #     use_safetensors=True, 
-        #     variant="fp16"
-        # )
-        # self.pipe_refiner.enable_model_cpu_offload()
+        # Load the refiner pipeline
+        self.pipe_refiner = DiffusionPipeline.from_pretrained(
+            self.model_refiner, 
+            torch_dtype=torch.float16, 
+            use_safetensors=True, 
+            variant="fp16"
+        )
+        self.pipe_refiner.enable_model_cpu_offload()
 
     def __call__(self, inputs):
         print("data",inputs)
@@ -38,39 +38,40 @@ class EndpointHandler:
 
         prompt, prompt2, negative_prompt, negative_prompt2 = inputs['prompt'], inputs['prompt2'], inputs['negative_prompt'], inputs['negative_prompt2']
         print(prompt, prompt2, negative_prompt, negative_prompt2)
-        # image_base_latent = self.pipe(
-        #                     prompt=prompt,
-        #                     prompt_2=prompt2,
-        #                     negative_prompt=negative_prompt,
-        #                     negative_prompt_2=negative_prompt2,
-        #                     guidance_scale=7.0,
-        #                     height=1024,
-        #                     width=1024,
-        #                     num_inference_steps=25,
-        #                     output_type="latent",
-        #                     denoising_end=0.8 # Cut the base denoising in 80%
-        #                 ).images[0]
+        
+        image_base_latent = self.pipe(
+                            prompt=prompt,
+                            prompt_2=prompt2,
+                            negative_prompt=negative_prompt,
+                            negative_prompt_2=negative_prompt2,
+                            guidance_scale=7.0,
+                            height=1024,
+                            width=1024,
+                            num_inference_steps=25,
+                            output_type="latent",
+                            denoising_end=0.8 # Cut the base denoising in 80%
+                        ).images[0]
 
-        # image_base_latent = image_base_latent.to("cuda")
+        image_base_latent = image_base_latent.to("cuda")
 
-        # # Refine the image
-        # image_refiner = self.pipe_refiner(
-        #                     prompt=prompt,
-        #                     prompt_2=prompt2,
-        #                     negative_prompt=negative_prompt,
-        #                     negative_prompt_2=negative_prompt2,
-        #                     image=image_base_latent,
-        #                     num_inference_steps=25,
-        #                     height=1024,
-        #                     width=1024,
-        #                     strength=0.3,
-        #                     denoising_start=0.8
-        #                 ).images[0]
+        # Refine the image
+        image_refiner = self.pipe_refiner(
+                            prompt=prompt,
+                            prompt_2=prompt2,
+                            negative_prompt=negative_prompt,
+                            negative_prompt_2=negative_prompt2,
+                            image=image_base_latent,
+                            num_inference_steps=25,
+                            height=1024,
+                            width=1024,
+                            strength=0.3,
+                            denoising_start=0.8
+                        ).images[0]
         
-        # # Convert the image to a format that can be easily outputted
-        # buffer = BytesIO()
-        # image_refiner.save(buffer, format="JPEG")
-        # buffer.seek(0)
-        # base64_encoded_result = base64.b64encode(buffer.read()).decode('utf-8')
+        # Convert the image to a format that can be easily outputted
+        buffer = BytesIO()
+        image_refiner.save(buffer, format="JPEG")
+        buffer.seek(0)
+        base64_encoded_result = base64.b64encode(buffer.read()).decode('utf-8')
         
         return {"image": "base64_encoded_result"}