|
from flask import url_for |
|
from diffusers import StableDiffusionPipeline |
|
import torch |
|
from fastapi import FastAPI, Response |
|
from fastapi.middleware.cors import CORSMiddleware |
|
from auth_token import auth_token |
|
|
|
app = FastAPI() |
|
|
|
app.add_middleware( |
|
CORSMiddleware, |
|
allow_credentials=True, |
|
allow_origins=["*"], |
|
allow_methods=["*"], |
|
allow_headers=["*"], |
|
) |
|
|
|
|
|
model_id = "runwayml/stable-diffusion-v1-5" |
|
pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=auth_token) |
|
pipe = pipe.to("cpu") |
|
pipe.enable_attention_slicing() |
|
|
|
def dummy(images, **kwargs): |
|
return images, False |
|
|
|
pipe.safety_checker = dummy |
|
|
|
@app.get("/") |
|
def hello(): |
|
return "Hello, I'm Artist" |
|
|
|
|
|
@app.get("/gen/{prompt}") |
|
def generate_image(prompt: str): |
|
image = pipe(prompt, |
|
guidance_scale=8.5 |
|
).images[0] |
|
|
|
image.save('image.png') |
|
|
|
|
|
return "Gen done" |
|
|