Spaces:
Paused
Paused
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()
}
|