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(fileUpload: UploadFile): logging.warning("in verify") logging.warning(fileUpload.filename) # check if the file has been uploaded if fileUpload.filename: # strip the leading path from the file name fn = os.path.basename(fileUpload.filename) # open read and write the file into the server open(fn, 'wb').write(fileUpload.file.read()) # subprocess.check_output( # [ # "./truepic", # "sign", # fileUpload.filename, # "--output", # (os.getcwd() + "/static/" + fileUpload.filename), # ] # ) subprocess.check_output( [ "export", "STEG_AI_API_KEY=" + os.environ.get("steg_api_key"), ] ) subprocess.check_output( [ "./scripts/sign.sh", fileUpload.filename, ] ) return {"response": fileUpload.filename} @app.post("/sign") def sign_image(fileUpload: UploadFile): logging.warning("in verify") logging.warning(fileUpload.filename) # check if the file has been uploaded if fileUpload.filename: # strip the leading path from the file name fn = os.path.basename(fileUpload.filename) # open read and write the file into the server open(fn, 'wb').write(fileUpload.file.read()) return {"response": fileUpload.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")