from fastapi import FastAPI from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse import subprocess import os import json import uuid import html import torch from diffusers import ( StableDiffusionPipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler, ) app = FastAPI() @app.get("/generate") def generate_image(prompt, model): torch.cuda.empty_cache() prompt = html.escape(prompt) model = html.escape(model) modelArray = model.split(",") modelName = modelArray[0] modelVersion = modelArray[1] pipeline = StableDiffusionPipeline.from_pretrained( str(modelName), torch_dtype=torch.float16 ) pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config) pipeline = pipeline.to("cuda") image = pipeline(prompt, num_inference_steps=50, height=512, width=512).images[0] filename = str(uuid.uuid4()) + ".jpg" image.save(filename) assertion = { "assertions": [ { "label": "com.truepic.custom.ai", "data": { "model_name": modelName, "model_version": modelVersion, "prompt": prompt, }, } ] } json_object = json.dumps(assertion) subprocess.check_output( [ "./truepic", "sign", filename, "--assertions-inline", json_object, "--output", (os.getcwd() + "/static/" + filename), ] ) return {"response": filename} app.mount("/", StaticFiles(directory="static", html=True), name="static") @app.get("/") def index() -> FileResponse: return FileResponse(path="/app/static/index.html", media_type="text/html")