|
from fastapi import FastAPI, Request, UploadFile, HTTPException, status |
|
from fastapi.staticfiles import StaticFiles |
|
from fastapi.responses import FileResponse |
|
import aiofiles |
|
|
|
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) |
|
|
|
|
|
if fileitem.filename: |
|
|
|
fn = os.path.basename(fileitem.filename) |
|
|
|
|
|
open(fn, 'wb').write(fileitem.file.read()) |
|
|
|
return {"response": fileitem.filename} |
|
|
|
@app.post('/upload') |
|
async def upload(fileUpload: UploadFile): |
|
try: |
|
contents = await fileUpload.read() |
|
async with aiofiles.open(fileUpload.filename, 'wb') as f: |
|
await f.write(contents) |
|
except Exception: |
|
raise HTTPException( |
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
|
detail='There was an error uploading the file', |
|
) |
|
finally: |
|
await fileUpload.close() |
|
|
|
return {'message': f'Successfuly uploaded {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") |
|
|
