add sd

main.py

@@ -1,7 +1,11 @@
+import os
+import time
+import sys
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from pydantic import Field
+from diffusers import StableDiffusionPipeline
 
 app = FastAPI()
 
@@ -20,6 +24,29 @@ def index():
 @app.post("/draw", response_model=ItemOut)
 def draw(data: Data):
 	if data.member_secret != "" and data.member_secret == os.environ.get("MEMBER_SECRET"):
-		return {"status": "OK", "file": "hello world!"}
+		print(f"Is CUDA available: {torch.cuda.is_available()}")
+
+		# prompt = '(('+data.string+')) (( photograph )), highly detailed, sharp focus, 8k, 4k, (( photorealism )), detailed, saturated, portrait, 50mm, F/2.8, 1m away, ( global illumination, studio light, volumetric light ), ((( multicolor lights )))'
+		prompt = '(('+data.string+')) (( photograph )), highly detailed, sharp focus, 8k, 4k, (( photorealism )), detailed, saturated, portrait, 50mm, F/2.8, 1m away, ((( multicolor lights )))'
+		n_prompt = 'text, blurry, art, painting, rendering, drawing, sketch, (( ugly )), (( duplicate )), ( morbid ), (( mutilated )), ( mutated ), ( deformed ), ( disfigured ), ( extra limbs ), ( malformed limbs ), ( missing arms ), ( missing legs ), ( extra arms ), ( extra legs ), ( fused fingers ), ( too many fingers ), long neck, low quality, worst quality'
+
+		# https://huggingface.co/docs/hub/spaces-sdks-docker-first-demo
+		# how to validation: https://qiita.com/bee2/items/75d9c0d7ba20e7a4a0e9
+		# https://github.com/huggingface/diffusers
+
+		model_id = 'stabilityai/stable-diffusion-2'
+
+		#pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)
+		pipe = StableDiffusionPipeline.from_pretrained(model_id)
+		#pipe = pipe.to('cuda')
+
+		image = pipe(prompt, negative_prompt=n_prompt).images[0]
+
+		fileName = "sd_" + str(time.time()) + '.png'
+		# image.save("/code/tmpdir/" + fileName)
+
+		print(fileName)
+
+		return {"status": "OK", "file": fileName}
 	else:
 		return {"status": "SORRY! This file is member only.", "file": ""}