import bentoml

from PIL.Image import Image

from annotated_types import Le, Ge

from typing_extensions import Annotated

from compel import Compel

from compel import Compel,ReturnedEmbeddingsType

from io import BytesIO

import base64

from <http://bentoml.io|bentoml.io> import JSON

MODEL_ID = "stabilityai/sdxl-turbo"

sample_prompt = "A cinematic shot of a baby racoon wearing an intricate italian priest robe."

@bentoml.service(

traffic={"timeout": 300},

workers=1,

resources={

"gpu": 1,

"gpu_type": "nvidia-l4",

},

)

class SDXLTurbo:

def __init__(self) -> None:

from diffusers import AutoPipelineForText2Image

import torch

self.pipe = AutoPipelineForText2Image.from_pretrained(

MODEL_ID,

torch_dtype=torch.float16,

variant="fp16",

use_safetensors=True,

)

<http://self.pipe.to|self.pipe.to>(device="cuda")

@bentoml.api

def txt2imgTest(

self,

prompt: str = sample_prompt,

negative_prompt: str = "",

num_inference_steps: Annotated[int, Ge(1), Le(10)] = 1,

guidance_scale: float = 0.0,

) -> Image:

image = self.pipe(prompt=prompt,

guidance_scale=guidance_scale,

num_inference_steps=1).images[0]

compel_proc = Compel(tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2],

text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2],

returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,

requires_pooled=[False, True])

conditioning, pooled = compel_proc(prompt)

neg_conditioning, neg_pooled = compel_proc(negative_prompt)

[conditioning, neg_conditioning] = compel_proc.pad_conditioning_tensors_to_same_length([conditioning, neg_conditioning])

image = self.pipe(prompt_embeds = conditioning, pooled_prompt_embeds = pooled,

negative_prompt_embeds = negative_prompt, negative_pooled_prompt_embeds = neg_pooled,

num_inference_steps = num_inference_steps).images[0]

return image