Save images locally. Add manual seed

.gitignore

@@ -1,2 +1,3 @@
 .env
-image.png
+image.png
+images/

docker-compose.yaml

@@ -10,10 +10,14 @@ services:
     env_file:
       - .env
     volumes:
-      - ~/Downloads/sd-v1-4-full-ema.ckpt:/app/model.ckpt
+      - ./images:/images
+      - ./notebooks:/notebooks
+      - ./models:/root/.cache/huggingface/diffusers/
+    ports:
+      - 8888:8888
+      - 8000:8000
 
   vosk:
-    image: alphacep/kaldi-en-gpu
+    image: alphacep/kaldi-en
     ports:
       - 2700:2700
-    runtime: nvidia

imageserver/main.py

@@ -9,6 +9,7 @@ from os import getenv
 from fastapi import FastAPI, Response, HTTPException
 from pydantic import BaseModel
 import io
+from PIL.PngImagePlugin import PngInfo
 
 load_dotenv()
 
@@ -30,16 +31,23 @@ app = FastAPI()
     },
     response_class=Response
 )
-def root(text: Text):
+def root(text: str):
-    prompt = text.text
+    prompt = text.replace('+', ' ')
     print(prompt)
     try:
-        image = pipe(prompt).images[0]
+        generator = torch.Generator("cuda").manual_seed(1024)
+        resp = pipe(prompt)
+        print(resp)
+        image = resp.images[0]
     except RuntimeError as e:
         raise HTTPException(status_code=202, detail="Busy")
-    except:     
+    except:
         raise HTTPException(status_code=504)
 
+    metadata = PngInfo()
+    metadata.add_text("text", prompt)
+    image.save(f'/images/{str(uuid.uuid4())}.png', pnginfo=metadata)    
+    
     imgByteArr = io.BytesIO()
     image.save(imgByteArr, format="PNG")
     imgByteArr = imgByteArr.getvalue()