how to create my API server for this model

#127
by LynchNancy - opened

my program code is as follows , but when I run this code ,it cannot give my right response , it will gives me" pic should be Tensor or ndarray. Got <class 'PIL.Image.Image'>." what's the problem , which master can help me ,thank you!!
from flask import Flask, request, make_response
from diffusers import DiffusionPipeline
import torch
from torchvision.transforms import ToPILImage
from io import BytesIO

app = Flask(name)

load both base & refiner

base = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
base.to("cuda")
refiner = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0",
text_encoder_2=base.text_encoder_2,
vae=base.vae,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
)
refiner.to("cuda")

@app .route('/predict', methods=['POST'])
def predict():
data = request.json
prompt = data['text']

Define how many steps and what % of steps to be run on each experts (80/20) here

n_steps = 40
high_noise_frac = 0.8

run both experts

image = base(
prompt=prompt,
num_inference_steps=n_steps,
denoising_end=high_noise_frac,
output_type="latent",
).images
image = refiner(
prompt=prompt,
num_inference_steps=n_steps,
denoising_start=high_noise_frac,
image=image,
).images[0]

Convert the Tensor to PIL Image

to_pil = ToPILImage()
pil_image = to_pil(image)

Convert the PIL Image to JPEG format

output = BytesIO()
pil_image.save(output, format='JPEG')
jpeg = output.getvalue()

Create a response object

resp = make_response(jpeg)

Set the response's headers

resp.headers['Content-Type'] = 'image/jpeg'

return resp

if name == 'main':
app.run(port=5000)

The simplest way is to use FastAPI
https://fastapi.tiangolo.com/

Sign up or log in to comment