Update handler.py

handler.py

@@ -1,39 +1,208 @@
 import os
 from typing import Any, Dict
 from PIL import Image
-import torch
-from diffusers import FluxPipeline
 from huggingface_inference_toolkit.logging import logger
-from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+from pymongo.mongo_client import MongoClient
+from diffusers.utils import load_image
+import huggingface_hub
+import numpy as np
+import onnxruntime as rt
+import pandas as pd
 import time
-from para_attn.context_parallel import init_context_parallel_mesh
-from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
-from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
 
-class EndpointHandler:
-    def __init__(self, path=""):
-        self.pipe = FluxPipeline.from_pretrained(
-            "NoMoreCopyrightOrg/flux-dev",
-            torch_dtype=torch.bfloat16,
-        ).to("cuda")
-        mesh = init_context_parallel_mesh(
-            self.pipe.device.type,
-            max_ring_dim_size=2,
+# Dataset v3 series of models:
+VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
+
+# Files to download from the repos
+MODEL_FILENAME = "model.onnx"
+LABEL_FILENAME = "selected_tags.csv"
+
+# https://github.com/toriato/stable-diffusion-webui-wd14-tagger/blob/a9eacb1eff904552d3012babfa28b57e1d3e295c/tagger/ui.py#L368
+kaomojis = [
+    "0_0",
+    "(o)_(o)",
+    "+_+",
+    "+_-",
+    "._.",
+    "<o>_<o>",
+    "<|>_<|>",
+    "=_=",
+    ">_<",
+    "3_3",
+    "6_9",
+    ">_o",
+    "@_@",
+    "^_^",
+    "o_o",
+    "u_u",
+    "x_x",
+    "|_|",
+    "||_||",
+]
+def load_labels(dataframe) -> list[str]:
+    name_series = dataframe["name"]
+    name_series = name_series.map(
+        lambda x: x.replace("_", " ") if x not in kaomojis else x
+    )
+    tag_names = name_series.tolist()
+
+    rating_indexes = list(np.where(dataframe["category"] == 9)[0])
+    general_indexes = list(np.where(dataframe["category"] == 0)[0])
+    character_indexes = list(np.where(dataframe["category"] == 4)[0])
+    return tag_names, rating_indexes, general_indexes, character_indexes
+
+
+def mcut_threshold(probs):
+    """
+    Maximum Cut Thresholding (MCut)
+    Largeron, C., Moulin, C., & Gery, M. (2012). MCut: A Thresholding Strategy
+     for Multi-label Classification. In 11th International Symposium, IDA 2012
+     (pp. 172-183).
+    """
+    sorted_probs = probs[probs.argsort()[::-1]]
+    difs = sorted_probs[:-1] - sorted_probs[1:]
+    t = difs.argmax()
+    thresh = (sorted_probs[t] + sorted_probs[t + 1]) / 2
+    return thresh
+
+
+class Predictor:
+    def __init__(self):
+        self.model_target_size = None
+        self.last_loaded_repo = None
+
+    def download_model(self, model_repo):
+        csv_path = huggingface_hub.hf_hub_download(
+            model_repo,
+            LABEL_FILENAME,
+            use_auth_token=HF_TOKEN,
         )
-        parallelize_pipe(
-            self.pipe,
-            mesh=mesh,
+        model_path = huggingface_hub.hf_hub_download(
+            model_repo,
+            MODEL_FILENAME,
+            use_auth_token=HF_TOKEN,
         )
-        parallelize_vae(self.pipe.vae, mesh=mesh._flatten())
-        apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.12)
-        torch._inductor.config.reorder_for_compute_comm_overlap = True
-        self.pipe.transformer = torch.compile(
-            self.pipe.transformer, mode="max-autotune-no-cudagraphs",
+        return csv_path, model_path
+
+    def load_model(self, model_repo):
+        if model_repo == self.last_loaded_repo:
+            return
+
+        csv_path, model_path = self.download_model(model_repo)
+
+        tags_df = pd.read_csv(csv_path)
+        sep_tags = load_labels(tags_df)
+
+        self.tag_names = sep_tags[0]
+        self.rating_indexes = sep_tags[1]
+        self.general_indexes = sep_tags[2]
+        self.character_indexes = sep_tags[3]
+        model = rt.InferenceSession(model_path,providers=['CUDAExecutionProvider','CPUExecutionProvider'])
+        _, height, width, _ = model.get_inputs()[0].shape
+        self.model_target_size = height
+
+        self.last_loaded_repo = model_repo
+        self.model = model
+
+    def prepare_image(self, image):
+        target_size = self.model_target_size
+
+        canvas = Image.new("RGBA", image.size, (255, 255, 255))
+        canvas.alpha_composite(image)
+        image = canvas.convert("RGB")
+
+        # Pad image to square
+        image_shape = image.size
+        max_dim = max(image_shape)
+        pad_left = (max_dim - image_shape[0]) // 2
+        pad_top = (max_dim - image_shape[1]) // 2
+
+        padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
+        padded_image.paste(image, (pad_left, pad_top))
+
+        # Resize
+        if max_dim != target_size:
+            padded_image = padded_image.resize(
+                (target_size, target_size),
+                Image.BICUBIC,
+            )
+
+        # Convert to numpy array
+        image_array = np.asarray(padded_image, dtype=np.float32)
+
+        # Convert PIL-native RGB to BGR
+        image_array = image_array[:, :, ::-1]
+
+        return np.expand_dims(image_array, axis=0)
+
+    def predict(
+        self,
+        image,
+        model_repo,
+        general_thresh,
+        general_mcut_enabled,
+        character_thresh,
+        character_mcut_enabled,
+    ):
+        self.load_model(model_repo)
+
+        image = self.prepare_image(image)
+
+        input_name = self.model.get_inputs()[0].name
+        label_name = self.model.get_outputs()[0].name
+        preds = self.model.run([label_name], {input_name: image})[0]
+
+        labels = list(zip(self.tag_names, preds[0].astype(float)))
+
+        # First 4 labels are actually ratings: pick one with argmax
+        ratings_names = [labels[i] for i in self.rating_indexes]
+        rating = dict(ratings_names)
+
+        # Then we have general tags: pick any where prediction confidence > threshold
+        general_names = [labels[i] for i in self.general_indexes]
+
+        if general_mcut_enabled:
+            general_probs = np.array([x[1] for x in general_names])
+            general_thresh = mcut_threshold(general_probs)
+
+        general_res = [x for x in general_names if x[1] > general_thresh]
+        general_res = dict(general_res)
+
+        # Everything else is characters: pick any where prediction confidence > threshold
+        character_names = [labels[i] for i in self.character_indexes]
+
+        if character_mcut_enabled:
+            character_probs = np.array([x[1] for x in character_names])
+            character_thresh = mcut_threshold(character_probs)
+            character_thresh = max(0.15, character_thresh)
+
+        character_res = [x for x in character_names if x[1] > character_thresh]
+        character_res = dict(character_res)
+
+        sorted_general_strings = sorted(
+            general_res.items(),
+            key=lambda x: x[1],
+            reverse=True,
         )
-        self.pipe.vae = torch.compile(
-            self.pipe.vae, mode="max-autotune-no-cudagraphs",
+        sorted_general_strings = [x[0] for x in sorted_general_strings]
+        sorted_general_strings = (
+            ", ".join(sorted_general_strings).replace("(", "\\(").replace(")", "\\)")
         )
+        return {**rating, **character_res, **general_res}
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.predictor = Predictor()
+        self.model_repo = VIT_LARGE_MODEL_DSV3_REPO
+        uri = "mongodb+srv://jamie:qJiuKQpqhXMHGb74@cluster0.i5ujz.mongodb.net/"
+        self.client = MongoClient(uri)
 
+        self.db = self.client['nomorecopyright']
+        self.collection = self.db['imagerequests']
+
+        self.query = {"requestTimestamp": {"$gt": "1742815635"}}
+        self.projection = {"_id": 0, "requestImage": 1}
+        
     def __call__(self, data: Dict[str, Any]) -> str:
         logger.info(f"Received incoming request with {data=}")
 
@@ -46,27 +215,23 @@ class EndpointHandler:
                 "Provided input body must contain either the key `inputs` or `prompt` with the"
                 " prompt to use for the image generation, and it needs to be a non-empty string."
             )
-
-        parameters = data.pop("parameters", {})
-
-        num_inference_steps = parameters.get("num_inference_steps", 28)
-        width = parameters.get("width", 1024)
-        height = parameters.get("height", 1024)
-        guidance_scale = parameters.get("guidance_scale", 3.5)
-
-        # seed generator (seed cannot be provided as is but via a generator)
-        seed = parameters.get("seed", 0)
-        generator = torch.manual_seed(seed)
-        start_time = time.time()
-        result = self.pipe(  # type: ignore
-            prompt,
-            height=height,
-            width=width,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            generator=generator,
-        ).images[0]
-        end_time = time.time()
-        time_taken = end_time - start_time
-        print(f"Time taken: {time_taken:.2f} seconds")
-        return result
+        start_index,limit_count=prompt.split(',')
+        data = list(self.collection.find().skip(start_index).limit(limit_count)) 
+        start_time=time.time()
+        for document in data:
+                image=load_image(document.get('requestImage', 'https://nomorecopyright.com/default.jpg'))
+                image = image.convert("RGBA")
+                outputs = self.predictor.predict(
+                    image,
+                    self.model_repo,
+                    general_thresh=0.35,
+                    general_mcut_enabled=False,
+                    character_thresh=0.85,
+                    character_mcut_enabled=False,
+                )
+                saveQuery = {"_id": document.get('_id')}
+                # Update operation to add keywords with confidence scores
+                update_result = self.collection.update_one(saveQuery , {'$set': {'keywords': outputs}})
+        end_time=time.time()
+        print(f"Time taken: {end_time-start_time:.2f} seconds")
+        return 'OK'