Spaces:
Sleeping
Sleeping
File size: 3,435 Bytes
7f871a4 7d54ba7 7f871a4 ee20381 8f77a4e 5aa8efd 8f77a4e ee20381 8f77a4e c92b67b ee20381 5aa8efd ee20381 7d54ba7 ee20381 7d54ba7 ee20381 7d54ba7 ee20381 7d54ba7 ee20381 7d54ba7 7f871a4 8e89c60 ee20381 c3efa02 8e89c60 7f871a4 6169018 8f77a4e ee20381 7d54ba7 ee20381 5aa8efd ee20381 7f871a4 ee20381 7f871a4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import numpy as np
import PIL.Image
import torch
from diffusers import LCMScheduler, AutoPipelineForText2Image
from fastapi import FastAPI
import uvicorn
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import RedirectResponse, StreamingResponse
import io
import os
from pathlib import Path
from db import Database
import uuid
import logging
from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from asyncio import Lock
logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
SPACE_ID = os.environ.get("SPACE_ID", "")
DEV = os.environ.get("DEV", "0") == "1"
DB_PATH = Path("/data/cache") if SPACE_ID else Path("./cache")
IMGS_PATH = DB_PATH / "imgs"
DB_PATH.mkdir(exist_ok=True, parents=True)
IMGS_PATH.mkdir(exist_ok=True, parents=True)
database = Database(DB_PATH)
generate_lock = Lock()
model_id = "segmind/Segmind-Vega"
adapter_id = "segmind/Segmind-VegaRT"
dtype = torch.bfloat16
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
pipe = AutoPipelineForText2Image.from_pretrained(
model_id, torch_dtype=torch.float16, variant="fp16"
)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")
pipe.load_lora_weights(adapter_id)
pipe.fuse_lora()
def generate(
prompt: str,
negative_prompt: str = "",
seed: int = 0,
) -> PIL.Image.Image:
generator = torch.Generator().manual_seed(seed)
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
generator=generator,
num_inference_steps=4,
guidance_scale=0,
).images[0]
return image
app = FastAPI()
origins = [
"https://huggingface.co",
"http://huggingface.co",
"https://huggingface.co/",
"http://huggingface.co/",
]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.middleware("http")
async def validate_origin(request: Request, call_next):
if DEV:
return await call_next(request)
if request.headers.get("referer") not in origins:
raise HTTPException(status_code=403, detail="Forbidden")
return await call_next(request)
@app.get("/image")
async def generate_image(
prompt: str, negative_prompt: str = "", seed: int = 2134213213
):
cached_img = database.check(prompt, negative_prompt, seed)
if cached_img:
logging.info(f"Image found in cache: {cached_img[0]}")
return StreamingResponse(open(cached_img[0], "rb"), media_type="image/jpeg")
logging.info(f"Image not found in cache, generating new image")
async with generate_lock:
pil_image = generate(prompt, negative_prompt, seed)
img_id = str(uuid.uuid4())
img_path = IMGS_PATH / f"{img_id}.jpg"
pil_image.save(img_path)
img_io = io.BytesIO()
pil_image.save(img_io, "JPEG")
img_io.seek(0)
database.insert(prompt, negative_prompt, str(img_path), seed)
return StreamingResponse(img_io, media_type="image/jpeg")
@app.get("/")
async def main():
# redirect to https://huggingface.co/spaces/multimodalart/stable-cascade
return RedirectResponse(
"https://multimodalart-stable-cascade.hf.space/?__theme=system"
)
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)
|