Feat: stable-diffusion-v2 (#2009)

* Upgrade stable diffusion model to v2



Signed-off-by: jagadeesh <jagadeeshj@ideas2it.com>
Co-authored-by: Mark Saroufim <marksaroufim@fb.com>
Co-authored-by: Ankith Gunapal <agunapal@meta.com>
Co-authored-by: Ankith Gunapal <agunapal@ischool.Berkeley.edu>
Co-authored-by: Mark Saroufim <marksaroufim@gmail.com>

examples/diffusers/Download_model.py

@@ -4,7 +4,7 @@
 TOKEN = "Token generated from Huggingface dashboard"
 
 pipeline = DiffusionPipeline.from_pretrained(
-    "CompVis/stable-diffusion-v1-4",
+    "stabilityai/stable-diffusion-2",
     revision="fp16",
     torch_dtype=torch.float16,
     use_auth_token=TOKEN,

examples/diffusers/Readme.md

@@ -4,6 +4,12 @@
 
 Set access token generated form Huggingface in `Download_model.py` file
 
+Install dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
 ```bash
 python Download_model.py
 ```

examples/diffusers/config.properties

@@ -7,4 +7,3 @@ install_py_dep_per_model=true
 load_models=all
 max_response_size=655350000
 model_store=/mnt/models/model_store
-model_snapshot={"name":"startup.cfg","modelCount":1,"models":{"stable-diffusion":{"1.0":{"defaultVersion":true,"marName":"stable-diffusion.mar","minWorkers":1,"maxWorkers":1,"batchSize":1,"maxBatchDelay":200,"responseTimeout":60}}}}

examples/diffusers/stable_diffusion_handler.py

@@ -5,7 +5,7 @@
 import diffusers
 import numpy as np
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import DiffusionPipeline
 
 from ts.torch_handler.base_handler import BaseHandler
 
@@ -42,7 +42,7 @@ def initialize(self, ctx):
         with zipfile.ZipFile(model_dir + "/model.zip", "r") as zip_ref:
             zip_ref.extractall(model_dir + "/model")
 
-        self.pipe = StableDiffusionPipeline.from_pretrained(model_dir + "/model")
+        self.pipe = DiffusionPipeline.from_pretrained(model_dir + "/model")
         self.pipe.to(self.device)
         logger.info("Diffusion model from path %s loaded successfully", model_dir)
 
@@ -76,7 +76,7 @@ def inference(self, inputs):
         """
         # Handling inference for sequence_classification.
         inferences = self.pipe(
-            inputs, guidance_scale=7.5, num_inference_steps=50
+            inputs, guidance_scale=7.5, num_inference_steps=50, height=768, width=768
         ).images
 
         logger.info("Generated image: '%s'", inferences)