File size: 3,209 Bytes
dfcd6b0
ce1dd07
 
 
dfcd6b0
a3614ce
50a1127
 
 
c64b66e
 
370c710
ffd715f
ce1dd07
50a1127
e5727cb
de8b3c6
5948e4d
 
 
 
4da3b9a
e5727cb
 
 
 
 
 
 
 
 
 
460cdbf
430abad
ce1dd07
 
 
c64b66e
 
e5727cb
c64b66e
 
 
370c710
 
c64b66e
 
5233da8
c64b66e
5948e4d
 
 
4da3b9a
c64b66e
 
370c710
e5727cb
 
c530969
 
e5727cb
 
0dfd0db
 
c64b66e
 
 
 
370c710
771c3c3
c64b66e
 
771c3c3
4da3b9a
5948e4d
c64b66e
 
 
278270a
7eb427a
50a1127
 
 
 
 
 
 
 
 
 
 
351b0b4
50a1127
 
 
 
 
4da3b9a
50a1127
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
from fastapi import APIRouter, Form, BackgroundTasks
from config import settings
import os
import json
import utils
import torch
import requests
from PIL import Image
from io import BytesIO
from pydantic import BaseModel
import base64
import uuid
from transformers import AutoTokenizer, AutoModelForCausalLM

from diffusers import StableDiffusionImg2ImgPipeline
from diffusers import StableDiffusionInpaintPipeline

# tokenizer = AutoTokenizer.from_pretrained("openlm-research/open_llama_7b")
# model = AutoModelForCausalLM.from_pretrained(
#     "openlm-research/open_llama_7b", device_map="auto", load_in_4bit=True
# )

# model_id_or_path = "runwayml/stable-diffusion-v1-5"
# pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)


pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    revision="fp16",
    torch_dtype=torch.float16,
)

pipe = pipe.to("cuda")


router = APIRouter()

class ActionBody(BaseModel):
    url: str
    maskUrl: str
    prompt: str
    strength: float
    guidance_scale: float
    resizeW: int
    resizeH: int

@router.post("/perform-action")
async def performAction(actionBody: ActionBody):
    
    # model_inputs = tokenizer(["A list of colors: red, blue"], return_tensors="pt").to("cuda")
    # generated_ids = model.generate(**model_inputs)
    # output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

    response = requests.get(actionBody.url)
    init_image = Image.open(BytesIO(response.content)).convert("RGB")
    init_image = init_image.resize((actionBody.resizeW, actionBody.resizeH))

    response = requests.get(actionBody.maskUrl)
    mask_image = Image.open(BytesIO(response.content)).convert("RGB")
    mask_image = mask_image.resize((actionBody.resizeW, actionBody.resizeH))

    
    # images = pipe(prompt=actionBody.prompt, image=init_image,  strength=actionBody.strength, guidance_scale=actionBody.guidance_scale).images
    images = pipe(prompt=actionBody.prompt, image=init_image,  mask_image=mask_image).images
    print(images)
    buffered = BytesIO()
    images[0].save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue())
    imgUUID = str(uuid.uuid4())
    images[0].save(imgUUID+".png")

    return {
        "imageName" : imgUUID+".png",
        "image": "data:image/jpeg;base64,"+img_str.decode(),
        # "output": output
    }


@router.get("/hi")
async def hifunction():
    
    url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
    response = requests.get(url)
    init_image = Image.open(BytesIO(response.content)).convert("RGB")
    init_image = init_image.resize((768, 512))
    prompt = "A fantasy landscape, trending on artstation"
    images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
    print(images)
    print(images[0])

    buffered = BytesIO()
    images[0].save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue())

    # images[0].save("fantasy_landscape.png")

    return {
        "image": "data:image/jpeg;base64,"+img_str.decode()
    }