File size: 2,299 Bytes
fd1dd45
e158b55
 
 
 
 
 
 
c2aa8fe
e158b55
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b16d75f
 
 
 
 
 
 
 
 
 
 
 
 
 
e158b55
 
 
b16d75f
e158b55
 
 
 
 
 
b16d75f
bb4cd3c
fd1dd45
4ae11c9
fd1dd45
 
 
 
 
 
b16d75f
fd1dd45
b16d75f
 
2da7508
b16d75f
c9a473f
fd1dd45
 
 
 
2da7508
 
fd1dd45
c2aa8fe
be4dbd7
e158b55
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
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,
        ]
    )

    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)

    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")