File size: 1,776 Bytes
59d59d3 |
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 |
import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetModel
from diffusers.pipelines import FluxControlNetPipeline
from PIL import Image
import io
class CustomHandler:
def __init__(self, model_dir):
# Load model and pipeline
self.controlnet = FluxControlNetModel.from_pretrained(
model_dir, torch_dtype=torch.bfloat16
)
self.pipe = FluxControlNetPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
controlnet=self.controlnet,
torch_dtype=torch.bfloat16
)
self.pipe.to("cuda")
def preprocess(self, data):
# Load image from file
image_file = data.get("control_image", None)
if not image_file:
raise ValueError("Missing control_image in input.")
image = Image.open(image_file)
w, h = image.size
# Upscale x4
return image.resize((w * 4, h * 4))
def postprocess(self, output):
# Save output image to a file-like object
buffer = io.BytesIO()
output.save(buffer, format="PNG")
buffer.seek(0) # Reset buffer pointer
return buffer
def inference(self, data):
# Preprocess input
control_image = self.preprocess(data)
# Generate output
output_image = self.pipe(
prompt=data.get("prompt", ""),
control_image=control_image,
controlnet_conditioning_scale=0.6,
num_inference_steps=28,
guidance_scale=3.5,
height=control_image.size[1],
width=control_image.size[0],
).images[0]
# Postprocess output
return self.postprocess(output_image)
|