from fastapi import FastAPI, UploadFile from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse import subprocess import os import json import uuid import logging import torch from diffusers import ( StableDiffusionPipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler, ) app = FastAPI() @app.get("/generate") def generate_image(prompt, model): torch.cuda.empty_cache() 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.post("/verify") def verify_image(file: UploadFile): logging.warning("in verify") logging.warning(file.filename) logging.warning('form') logging.warning(form) logging.warning('fileitem') logging.warning(fileitem) # check if the file has been uploaded if fileitem.filename: # strip the leading path from the file name fn = os.path.basename(fileitem.filename) # open read and write the file into the server open(fn, 'wb').write(fileitem.file.read()) return {"response": fileitem.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")