how to create my API server for this model
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/