Merge pull request #1 from 0xb79/main

Update code to use data folder and clear cache

Author: ecmulli

README.MD

@@ -7,6 +7,9 @@
     * `python3 -m venv venv`
     * `source venv/bin/activate`
 2. Install requirements `pip install -r requirements.txt` 
+3. Sign up for Hugging Face and generate auth token (https://huggingface.co/docs/hub/security-tokens)
+4. Request access to stable diffusion model https://huggingface.co/CompVis/stable-diffusion-v1-4
+5. Input access token to `access_token` in `server.py` and `stable-diffusion.py`
 
 ## Playground
 
@@ -16,4 +19,4 @@
 ## Server
 
 1. run `venv/bin/python server.py`
-2. endpoint available at `<ipaddress>:port/stable-diffusion/?prompt=<prompt>` 
+2. endpoint available at `<ipaddress>:port/stable-diffusion?prompt=<prompt>` 

server.py

@@ -1,8 +1,13 @@
 from flask import Flask, request, send_file
 from torch import autocast, cuda
+import torch, os
 from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
 app = Flask(__name__)
 
+access_token = "enter access token"
+
+
+    
 # this will substitute the default PNDM scheduler for K-LMS  
 lms = LMSDiscreteScheduler(
     beta_start=0.00085, 
@@ -13,7 +18,8 @@
 pipe = StableDiffusionPipeline.from_pretrained(
     "CompVis/stable-diffusion-v1-4", 
     scheduler=lms,
-    use_auth_token=True
+    torch_dtype=torch.float16,
+    use_auth_token=access_token
 )
 
 if cuda.is_available():
@@ -23,7 +29,7 @@
 @app.route("/stable_diffusion")
 def get_result():
     prompt = request.args.get("prompt")
-    prompt_path = "data/" + prompt.replace(" ", "_") + ".png"
+    prompt_path = os.path.join(root_dir(), "data",prompt.replace(" ", "_") + ".png")
     print(prompt, prompt_path)
     if cuda.is_available():
         with autocast("cuda"):
@@ -32,8 +38,12 @@ def get_result():
         image = pipe(prompt)["sample"][0] 
 
     image.save(prompt_path)
-
+    torch.cuda.empty_cache()
+    
     return send_file(prompt_path, mimetype='image/png')
 
+def root_dir():
+    return os.path.abspath(os.path.dirname(__file__))
+    
 if __name__=="__main__":
-    app.run(debug = True)
+    app.run(host="127.0.0.1", port="5555",debug = True)

stable-diffusion.py

@@ -1,6 +1,9 @@
 from torch import autocast
 from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
 from IPython.display import Image
+import os
+
+access_token = "enter access token"
 
 # this will substitute the default PNDM scheduler for K-LMS  
 lms = LMSDiscreteScheduler(
@@ -12,13 +15,16 @@
 pipe = StableDiffusionPipeline.from_pretrained(
     "CompVis/stable-diffusion-v1-4", 
     scheduler=lms,
-    use_auth_token=True
+    use_auth_token=access_token
 ).to("cuda")
 
 prompt = "a photo of an astronaut riding a horse on mars"
-file_name = "astronaut_rides_horse.png"
+file_name = os.path.join(root_dir(), "data", "astronaut_rides_horse.png")
 
 with autocast("cuda"):
     image = pipe(prompt)["sample"][0]  
+    image.save(file_name)
 
-image.save(file_name)
+torch.cuda.empty_cache()
+def root_dir():
+    return os.path.abspath(os.path.dirname(__file__))