update handler

handler.py

@@ -1,56 +1,78 @@
-from diffusers import DiffusionPipeline
-from diffusers.utils import load_image
-from PIL import Image
-import torch
 import base64
 from io import BytesIO
-import sys
-import os
-from pathlib import Path
-
-model_dir = Path(__file__).parent.resolve()  # 获取handler.py所在目录的绝对路径
+from typing import Any, Dict
 
+import torch
+from diffusers import FluxFillPipeline
+from PIL import Image
 
-# 关键：注册 FluxFillPipeline 类（自动从 model_index.json 解析）
-pipe = DiffusionPipeline.from_pretrained(
-    str(model_dir),
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-).to("cuda" if torch.cuda.is_available() else "cpu")
 
 def decode_image(b64_string):
     image_data = base64.b64decode(b64_string)
     return Image.open(BytesIO(image_data)).convert("RGB")
 
+
 def encode_image(image):
     buffer = BytesIO()
     image.save(buffer, format="PNG")
     return base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-def handler(data):
-    try:
-        inputs = data.get("inputs", {})
-        prompt = inputs.get("prompt", "写实风格形象照")
-        image_b64 = inputs.get("image")
-        mask_b64 = inputs.get("mask")
 
-        if not image_b64 or not mask_b64:
-            return {"error": "缺少 image 或 mask 参数", "status": "failed"}
+class EndpointHandler:
+    def __init__(self, path="shangguanyanyan/flux1-fill-dev-custom"):
+        self.pipe = FluxFillPipeline.from_pretrained(
+            path, torch_dtype=torch.bfloat16
+        ).to("cuda" if torch.cuda.is_available() else "cpu")
 
-        image = decode_image(image_b64)
-        mask = decode_image(mask_b64)
+        self.parameters = {
+            "height": 1632,
+            "width": 1232,
+            "guidance_scale": 30,
+            "num_inference_steps": 50,
+            "max_sequence_length": 512,
+            "generator": torch.Generator("cpu").manual_seed(0),
+        }
 
-        result = pipe(
+    def __call__(self, data: Any) -> Dict[str, Any]:
+        """
+        data: {
+            "inputs": {
+                "image": base64_image,
+                "mask": base64_mask,
+                "prompt": prompt
+            },
+            "parameters": {
+                "height": 1632,
+                "width": 1232,
+                "guidance_scale": 30,
+                "num_inference_steps": 50,
+                "max_sequence_length": 512,
+            }
+        }
+        """
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
+
+        parameters.update(self.parameters)
+        base64_image = inputs.pop("image", "")
+        base64_mask = inputs.pop("mask", "")
+        prompt = inputs.pop("prompt", "")
+
+        if not base64_image or not base64_mask or not prompt:
+            return {
+                "error": "Please provide image, mask and prompt",
+                "status": "failed",
+            }
+
+        image = decode_image(base64_image)
+        mask = decode_image(base64_mask)
+
+        image = self.pipe(
             prompt=prompt,
             image=image,
             mask_image=mask,
-            height=int(inputs.get("height", 1024)),
-            width=int(inputs.get("width", 1024)),
-            guidance_scale=float(inputs.get("guidance_scale", 7.5)),
-            num_inference_steps=int(inputs.get("steps", 30))
+            **parameters,
         ).images[0]
 
-        return {"image": encode_image(result), "status": "success"}
-
-    except Exception as e:
-        return {"error": str(e), "status": "failed"}
+        return {"image": encode_image(image), "status": "success"}
 

local_test.py

@@ -0,0 +1,38 @@
+from diffusers.utils import load_image
+
+from handler import EndpointHandler, encode_image
+
+# init handler
+my_handler = EndpointHandler(path=".")
+# my_handler = EndpointHandler(path="black-forest-labs/FLUX.1-Fill-dev")
+
+image = load_image(
+    "./cup.png"
+)
+mask = load_image(
+    "./cup_mask.png"
+)
+prompt = "a white paper cup"
+
+data = {
+    "inputs": {
+        "image": encode_image(image),
+        "mask": encode_image(mask),
+        "prompt": prompt,
+    },
+    "parameters": {
+        "height": 1632,
+        "width": 1232,
+        "guidance_scale": 30,
+        "num_inference_steps": 50,
+        "max_sequence_length": 512,
+    },
+}
+# test the handler
+print("out")
+exit()
+result = my_handler(data=data)
+
+
+print("result:", result)
+

model_index.json

@@ -1,14 +1,32 @@
 {
   "_class_name": "FluxFillPipeline",
   "_diffusers_version": "0.32.0.dev0",
-  "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
-  "text_encoder": ["transformers", "CLIPTextModel", { "weight_name": "model.safetensors" }],
-  "text_encoder_2": ["transformers", "T5EncoderModel", { "weight_name": "model-00001-of-00002.safetensors" }],
-  "tokenizer": ["transformers", "CLIPTokenizer"],
-  "tokenizer_2": ["transformers", "T5TokenizerFast"],
-  "transformer": ["diffusers", "FluxTransformer2DModel", { "weight_name": "diffusion_pytorch_model.safetensors" }],
-  "vae": ["diffusers", "AutoencoderKL", { "weight_name": "diffusion_pytorch_model.safetensors" }],
-  "image_encoder": ["transformers", "CLIPVisionModelWithProjection", { "weight_name": "model.safetensors" }],
-  "unet": ["diffusers", "UNet2DConditionModel", { "weight_name": "diffusion_pytorch_model.safetensors" }]
-}
-
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModel"
+  ],
+  "text_encoder_2": [
+    "transformers",
+    "T5EncoderModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "tokenizer_2": [
+    "transformers",
+    "T5TokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "FluxTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

requirements.txt

@@ -3,3 +3,5 @@ transformers
 torch
 accelerate
 safetensors
+protobuf
+sentencepiece