import os import time import sys import re import random import torch from fastapi import FastAPI from fastapi.staticfiles import StaticFiles from pydantic import BaseModel from pydantic import Field from diffusers import StableDiffusionPipeline app = FastAPI() class Data(BaseModel): string: str member_secret: str class ItemOut(BaseModel): status: str file: str @app.get("/") def index(): return "SORRY! This file is member only." @app.post("/draw", response_model=ItemOut) def draw(data: Data): if data.member_secret != "" and data.member_secret == os.environ.get("MEMBER_SECRET"): print(f"Is CUDA available: {torch.cuda.is_available()}") if '_seed' in data.string: seed = 1024 else: seed = random.randrange(1024) text = re.sub('^#', '', data.string) text = re.sub('_seed', '', text) # prompt = '(('+text+')) (( 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 = '(('+text+')) (( 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) pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16) pipe = pipe.to('cuda') generator = torch.Generator("cuda").manual_seed(seed) image = pipe(prompt, negative_prompt=n_prompt, guidance_scale=7.5, generator=generator).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": ""} app.mount("/static", StaticFiles(directory="/code/tmpdir"), name="/static")