Update stable_diffusion_handler.py

stable_diffusion_handler.py

@@ -3,7 +3,7 @@ from abc import ABC
 
 import diffusers
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionXLPipeline
 
 from ts.torch_handler.base_handler import BaseHandler
 import numpy as np
@@ -41,7 +41,7 @@ class DiffusersHandler(BaseHandler, ABC):
         )
 
         
-        self.pipe = StableDiffusionPipeline.from_pretrained("./")
+        self.pipe = StableDiffusionXLPipeline.from_pretrained("./")
         self.pipe.to(self.device)
         logger.info("Diffusion model from path %s loaded successfully", model_dir)