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( [ "./scripts/sign.sh", filename, ] ) subprocess.check_output( [ "cp", "output.jpg", "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()) response = subprocess.check_output( [ "./scripts/verify.sh", fileUpload.filename, ] ) logging.warning(response) response_list = response.splitlines() c2pa_string = str(response_list[0]) c2pa = c2pa_string.split(":", 1) c2pa = c2pa[1].translate(None, " '") watermark_string = str(response_list[1]) watermark = watermark_string.split(":", 1) watermark = watermark[1].translate(None, " '") original_media_string = str(response_list[2]) original_media = original_media_string.split(":", 1) original_media = original_media[1].translate(None, " '") return {"response": fileUpload.filename, "contains_c2pa" : c2pa, "contains_watermark" : watermark, "original_media" : original_media} 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")