Spaces:
Sleeping
Sleeping
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 | |
def index(): | |
return "SORRY! This file is member only." | |
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") | |