update handler.py

handler.py

@@ -4,13 +4,18 @@ import torch
 import base64
 from io import BytesIO
 
+# ✅ 使用 huggingface repo id 形式加载（确保路径正确）
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    "./", torch_dtype=torch.float16
+    "shangguanyanyan/flux1-fill-dev-custom",  # 请确认仓库ID无误
+    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")
+    try:
+        image_data = base64.b64decode(b64_string)
+        return Image.open(BytesIO(image_data)).convert("RGB")
+    except Exception as e:
+        raise ValueError(f"解码图像失败: {str(e)}")
 
 def encode_image(image):
     buffer = BytesIO()
@@ -18,13 +23,47 @@ def encode_image(image):
     return base64.b64encode(buffer.getvalue()).decode("utf-8")
 
 def handler(data):
-    inputs = data.get("inputs", {})
-    prompt = inputs.get("prompt", "商务风格形象照，高清写实，蓝色西装")
-    image_b64 = inputs.get("image")
-    mask_b64 = inputs.get("mask")
+    try:
+        inputs = data.get("inputs", {})
+        prompt = inputs.get("prompt", "高清写实风格人物形象照")
 
-    image = decode_image(image_b64)
-    mask = decode_image(mask_b64)
+        image_b64 = inputs.get("image")
+        mask_b64 = inputs.get("mask")
+        if not image_b64 or not mask_b64:
+            raise ValueError("缺少必要的 image 或 mask 参数")
 
-    result = pipe(prompt=prompt, image=image, mask_image=mask).images[0]
-    return {"image": encode_image(result)}
+        image = decode_image(image_b64)
+        mask = decode_image(mask_b64)
+
+        # 默认参数（支持调整）
+        height = int(inputs.get("height", image.height))
+        width = int(inputs.get("width", image.width))
+        steps = int(inputs.get("num_inference_steps", 30))
+        cfg_scale = float(inputs.get("guidance_scale", 7.5))
+
+        image = image.resize((width, height))
+        mask = mask.resize((width, height))
+
+        result = pipe(
+            prompt=prompt,
+            image=image,
+            mask_image=mask,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale
+        ).images[0]
+
+        return {
+            "status": "success",
+            "image": encode_image(result),
+            "meta": {
+                "prompt": prompt,
+                "size": f"{width}x{height}",
+                "steps": steps,
+                "cfg_scale": cfg_scale
+            }
+        }
+
+    except Exception as e:
+        return {"status": "failed", "error": str(e)}