File size: 2,998 Bytes
fd1dd45
e158b55
 
 
 
 
 
 
c2aa8fe
e158b55
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bb4cd3c
fd1dd45
4ae11c9
fd1dd45
0e5ab33
fd1dd45
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c2aa8fe
fd1dd45
 
 
 
 
 
c2aa8fe
e158b55
fd1dd45
e158b55
fd1dd45
e158b55
 
fd1dd45
0e5ab33
fd1dd45
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
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),
            ]
        )

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