Merge pull request #42 from tushar2407/sd2.1

feat: stable diffusion 2.1 support
stochasticai · Jun 14, 2023 · dcee5d1 · dcee5d1
2 parents 8fd180d + f53ede4
commit dcee5d1
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/PyTorch/README.md b/PyTorch/README.md
@@ -29,7 +29,7 @@ docker build --build-arg model_dir_path=/path/to/stable_diffusion/model -f Docke
 ```
 
 2. Execute the Docker Container
-```
+```bash
 sudo docker run --gpus all -p 5000:5000 stable_diffusion_img
 ```
 
@@ -43,7 +43,7 @@ uvicorn server:app --host 0.0.0.0 --port 5000
 ## 3. Demo App
 
 To generate images as a command line tool, execute the following command:
-```
+```bash
 python demo.py --prompt "an astronaut riding a horse"
 ```
 

diff --git a/PyTorch/model.py b/PyTorch/model.py
@@ -5,7 +5,7 @@
 
 
 def load_model(
-    model_name_or_path="CompVis/stable-diffusion-v1-4",
+    model_name_or_path="stabilityai/stable-diffusion-2-1",
 ) -> StableDiffusionPipeline:
     """Load model
 
@@ -14,9 +14,9 @@ def load_model(
     """
     pipe = StableDiffusionPipeline.from_pretrained(
         model_name_or_path,
-        revision="fp16",
+        # revision="fp16",
         torch_dtype=torch.float16,
-        use_auth_token=True,
+        # use_auth_token=True,
     )
     pipe = pipe.to("cuda")
 

diff --git a/PyTorch/requirements.txt b/PyTorch/requirements.txt
@@ -1,5 +1,9 @@
-transformers==4.21.2
+transformers==4.30.1
 diffusers==0.3.0
-torch==1.12.1+cu116
+torch==2.0.1+cu117
 fastapi==0.85.0
-uvicorn[standard]==0.18.3
+uvicorn[standard]==0.18.3
+accelerate==0.20.3
+safetensors==0.3.1
+scipy==1.10.1
+torchvision==0.15.2+cu117