Upload 3 files

config.json

@@ -0,0 +1,21 @@
+{
+    "_class_name": "FluxControlNetModel",
+    "_diffusers_version": "0.31.0.dev0",
+    "_name_or_path": "/data/checkpoints/flux_controlnet_hf/controlnet_upscaling//",
+    "attention_head_dim": 128,
+    "axes_dims_rope": [
+      16,
+      56,
+      56
+    ],
+    "guidance_embeds": true,
+    "in_channels": 64,
+    "joint_attention_dim": 4096,
+    "num_attention_heads": 24,
+    "num_layers": 5,
+    "num_mode": null,
+    "num_single_layers": 0,
+    "patch_size": 1,
+    "pooled_projection_dim": 768
+  }
+  

diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a7ea24d2037ff2aa4d25f8b4ce9fe7e739a2cfe6b9d05106788005d5058c8ca
+size 3583232168

handler.py

@@ -0,0 +1,52 @@
+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)