PD model and funcitonnal endpoint inference + check progress'

feature_extractor/preprocessor_config.json

@@ -14,7 +14,7 @@
     0.4578275,
     0.40821073
   ],
-  "image_processor_type": "CLIPImageProcessor",
+  "image_processor_type": "CLIPFeatureExtractor",
   "image_std": [
     0.26862954,
     0.26130258,

handler.py

@@ -7,6 +7,9 @@ from pprint import pprint
 from typing import Any, Dict, List
 import os
 from pathlib import Path
+from typing import Union
+from concurrent.futures import ThreadPoolExecutor
+import numpy as np
 
 import torch
 from diffusers import (
@@ -14,12 +17,12 @@ from diffusers import (
     DPMSolverMultistepScheduler,
     DPMSolverSinglestepScheduler,
     EulerAncestralDiscreteScheduler,
+    utils,
 )
 from safetensors.torch import load_file
-from torch import autocast
-
-# https://huggingface.co/philschmid/stable-diffusion-v1-4-endpoints
-# https://huggingface.co/docs/inference-endpoints/guides/custom_handler
+from torch import autocast, tensor
+import torchvision.transforms
+from PIL import Image
 
 REPO_DIR = Path(__file__).resolve().parent
 
@@ -48,6 +51,7 @@ class EndpointHandler:
         "detailed_eye-10": str(REPO_DIR / "lora/detailed_eye-10.safetensors"),
         "add_detail": str(REPO_DIR / "lora/add_detail.safetensors"),
         "MuscleGirl_v1": str(REPO_DIR / "lora/MuscleGirl_v1.safetensors"),
+        "flat2": str(REPO_DIR / "lora/flat2.safetensors"),
     }
 
     TEXTUAL_INVERSION = [
@@ -55,10 +59,6 @@ class EndpointHandler:
             "weight_name": str(REPO_DIR / "embeddings/EasyNegative.safetensors"),
             "token": "easynegative",
         },
-        {
-            "weight_name": str(REPO_DIR / "embeddings/EasyNegative.safetensors"),
-            "token": "EasyNegative",
-        },
         {
             "weight_name": str(REPO_DIR / "embeddings/badhandv4.pt"),
             "token": "badhandv4",
@@ -69,16 +69,12 @@ class EndpointHandler:
         },
         {
             "weight_name": str(REPO_DIR / "embeddings/NegfeetV2.pt"),
-            "token": "NegfeetV2",
+            "token": "negfeetv2",
         },
         {
             "weight_name": str(REPO_DIR / "embeddings/ng_deepnegative_v1_75t.pt"),
             "token": "ng_deepnegative_v1_75t",
         },
-        {
-            "weight_name": str(REPO_DIR / "embeddings/ng_deepnegative_v1_75t.pt"),
-            "token": "NG_DeepNegative_V1_75T",
-        },
         {
             "weight_name": str(REPO_DIR / "embeddings/bad-hands-5.pt"),
             "token": "bad-hands-5",
@@ -86,6 +82,15 @@ class EndpointHandler:
     ]
 
     def __init__(self, path="."):
+        self.inference_progress = {}  # Dictionary to store progress of each request
+        self.inference_images = {}  # Dictionary to store latest image of each request
+        self.total_steps = {}
+        self.inference_in_progress = False
+
+        self.executor = ThreadPoolExecutor(
+            max_workers=1
+        )  # Vous pouvez ajuster max_workers en fonction de vos besoins
+
         # load the optimized model
         self.pipe = DiffusionPipeline.from_pretrained(
             path,
@@ -94,30 +99,31 @@ class EndpointHandler:
         )
         self.pipe = self.pipe.to(device)
 
+        # https://stablediffusionapi.com/docs/a1111schedulers/
+
         # DPM++ 2M SDE Karras
         # increase step to avoid high contrast num_inference_steps=30
+        # self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+        #     self.pipe.scheduler.config,
+        #     use_karras_sigmas=True,
+        #     algorithm_type="sde-dpmsolver++",
+        # )
+        # DPM++ 2M Karras
+        # increase step to avoid high contrast num_inference_steps=30
         self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
             self.pipe.scheduler.config,
             use_karras_sigmas=True,
-            algorithm_type="sde-dpmsolver++",
         )
 
         # Mode boulardus
         self.pipe.safety_checker = None
 
+        # Disable progress bar
+        self.pipe.set_progress_bar_config(disable=True)
+
         # Load negative embeddings to avoid bad hands, etc
         self.load_embeddings()
 
-        # Load default Lora models
-        self.pipe = self.load_selected_loras(
-            [
-                ("polyhedron_new_skin_v1.1", 0.35),  # nice Skin
-                ("detailed_eye-10", 0.3),  # nice eyes
-                ("add_detail", 0.4),  # detailed pictures
-                ("MuscleGirl_v1", 0.3),  # shape persons
-            ],
-        )
-
         # boosts performance by another 20%
         self.pipe.enable_xformers_memory_efficient_attention()
         self.pipe.enable_attention_slicing()
@@ -215,14 +221,121 @@ class EndpointHandler:
             )
         return self.pipe
 
-    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
-        """
-        Args:
-            data (:obj:):
-                includes the input data and the parameters for the inference.
-        Return:
-            A :obj:`dict`:. base64 encoded image
-        """
+    def __call__(self, data: Any) -> Dict:
+        """Handle incoming requests."""
+
+        action = data.get("action", None)
+        request_id = data.get("request_id")
+
+        # Check if the request_id is valid for all actions
+        if not request_id:
+            return {"flag": "error", "message": "Missing request_id."}
+
+        if action == "check_progress":
+            return self.check_progress(request_id)
+
+        elif action == "inference":
+            # Check if an inference is already in progress
+            if self.inference_in_progress:
+                return {
+                    "flag": "error",
+                    "message": "Another inference is already in progress. Please wait.",
+                }
+
+            # Set the inference state to in progress
+            self.clean_request_data(request_id)
+            self.inference_in_progress = True
+            self.inference_progress[request_id] = 0
+            self.inference_images[request_id] = None
+
+            self.executor.submit(self.start_inference, data)
+
+            return {
+                "flag": "success",
+                "message": "Inference started",
+                "request_id": request_id,
+            }
+
+        else:
+            return {"flag": "error", "message": f"Unsupported action: {action}"}
+
+    def clean_request_data(self, request_id: str):
+        """Clean up the data related to a specific request ID."""
+
+        # Remove the request ID from the progress dictionary
+        self.inference_progress.pop(request_id, None)
+
+        # Remove the request ID from the images dictionary
+        self.inference_images.pop(request_id, None)
+
+        # Remove the request ID from the total_steps dictionary
+        self.total_steps.pop(request_id, None)
+
+        # Set inference to False
+        self.inference_in_progress = False
+
+    def progress_callback(
+        self,
+        step: int,
+        timestep: int,
+        latents: Any,
+        request_id: str,
+        status: str,
+    ):
+        try:
+            if status == "progress":
+                # Latents to numpy
+                img_data = self.pipe.decode_latents(latents)
+                img_data = (img_data.squeeze() * 255).astype(np.uint8)
+                img = Image.fromarray(img_data, "RGB")
+                # print(img_data)
+            else:
+                # pil object
+                # print(latents)
+                img = latents
+
+            buffered = BytesIO()
+            img.save(buffered, format="PNG")
+
+            # print(status)
+            # Save the image to a file
+            # img.save("squirel.png", format="PNG")
+
+            # Encode the image into a base64 string representation
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+
+        except Exception as e:
+            print(f"Error: {e}")
+
+        # Store progress and image
+        progress_percentage = (
+            step / self.total_steps[request_id]
+        ) * 100  # Assuming self.total_steps is the total number of steps for inference
+
+        self.inference_progress[request_id] = progress_percentage
+        self.inference_images[request_id] = img_str
+
+    def check_progress(self, request_id: str) -> Dict[str, Union[str, float]]:
+        progress = self.inference_progress.get(request_id, 0)
+        latest_image = self.inference_images.get(request_id, None)
+
+        # print(self.inference_progress)
+
+        if progress >= 100:
+            status = "complete"
+        else:
+            status = "in-progress"
+
+        return {
+            "flag": "success",
+            "status": status,
+            "progress": int(progress),
+            "image": latest_image,
+        }
+
+    def start_inference(self, data: Dict) -> Dict:
+        """Start a new inference."""
+
         global device
 
         # Which Lora do we load ?
@@ -241,8 +354,8 @@ class EndpointHandler:
             "width",
             "num_inference_steps",
             "height",
-            "seed",
             "guidance_scale",
+            "request_id",
         ]
 
         missing_fields = [field for field in required_fields if field not in data]
@@ -256,17 +369,21 @@ class EndpointHandler:
         # Now extract the fields
         prompt = data["prompt"]
         negative_prompt = data["negative_prompt"]
-        loras_model = data.pop("loras_model", None)
-        seed = data["seed"]
+        loras_model = data.get("loras_model", None)
+        seed = data.get("seed", None)
         width = data["width"]
         num_inference_steps = data["num_inference_steps"]
         height = data["height"]
         guidance_scale = data["guidance_scale"]
+        request_id = data["request_id"]
+
+        # Used for progress checker
+        self.total_steps[request_id] = num_inference_steps
 
         # USe this to add automatically some negative prompts
         forced_negative = (
             negative_prompt
-            + """easynegative, badhandv4, bad-artist-anime, NegfeetV2, ng_deepnegative_v1_75t, bad-hands-5  """
+            + """, easynegative, badhandv4, bad-artist-anime, negfeetv2, ng_deepnegative_v1_75t, bad-hands-5, """
         )
 
         # Set the generator seed if provided
@@ -288,15 +405,20 @@ class EndpointHandler:
                     negative_prompt=forced_negative,
                     generator=generator,
                     max_embeddings_multiples=5,
+                    callback=lambda step, timestep, latents: self.progress_callback(
+                        step, timestep, latents, request_id, "progress"
+                    ),
+                    callback_steps=8,  #  The frequency at which the callback function is called.
+                    # output_type="pt",
                 ).images[0]
 
-            # encode image as base 64
-            buffered = BytesIO()
-            image.save(buffered, format="JPEG")
-            img_str = base64.b64encode(buffered.getvalue())
+            # print(image)
+            self.progress_callback(
+                num_inference_steps, 0, image, request_id, "complete"
+            )
 
-            # Return the success response
-            return {"flag": "success", "image": img_str.decode()}
+            # for debug
+            # image.save("squirelb.png", format="PNG")
 
         except Exception as e:
             # Handle any other exceptions and return an error response

lora/flat2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:054e950e72181bb45ddbc7106d3625de406477725b5b313a91fe4522f03dbe0a
+size 6865699

model_index.json

@@ -3,7 +3,7 @@
   "_diffusers_version": "0.20.0",
   "feature_extractor": [
     "transformers",
-    "CLIPImageProcessor"
+    "CLIPFeatureExtractor"
   ],
   "requires_safety_checker": true,
   "safety_checker": [

safety_checker/config.json

@@ -15,7 +15,7 @@
     "attention_dropout": 0.0,
     "bad_words_ids": null,
     "begin_suppress_tokens": null,
-    "bos_token_id": 0,
+    "bos_token_id": 49406,
     "chunk_size_feed_forward": 0,
     "cross_attention_hidden_size": null,
     "decoder_start_token_id": null,
@@ -24,7 +24,7 @@
     "dropout": 0.0,
     "early_stopping": false,
     "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": 2,
+    "eos_token_id": 49407,
     "exponential_decay_length_penalty": null,
     "finetuning_task": null,
     "forced_bos_token_id": null,
@@ -80,17 +80,11 @@
     "top_p": 1.0,
     "torch_dtype": null,
     "torchscript": false,
-    "transformers_version": "4.25.1",
+    "transformers_version": "4.31.0",
     "typical_p": 1.0,
     "use_bfloat16": false,
     "vocab_size": 49408
   },
-  "text_config_dict": {
-    "hidden_size": 768,
-    "intermediate_size": 3072,
-    "num_attention_heads": 12,
-    "num_hidden_layers": 12
-  },
   "torch_dtype": "float32",
   "transformers_version": null,
   "vision_config": {
@@ -167,15 +161,8 @@
     "top_p": 1.0,
     "torch_dtype": null,
     "torchscript": false,
-    "transformers_version": "4.25.1",
+    "transformers_version": "4.31.0",
     "typical_p": 1.0,
     "use_bfloat16": false
-  },
-  "vision_config_dict": {
-    "hidden_size": 1024,
-    "intermediate_size": 4096,
-    "num_attention_heads": 16,
-    "num_hidden_layers": 24,
-    "patch_size": 14
   }
 }

safety_checker/pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:16d28f2b37109f222cdc33620fdd262102ac32112be0352a7f77e9614b35a394
-size 1216064769
+oid sha256:753acd54aa6d288d6c0ce9d51468eb28f495fcbaacf0edf755fa5fc7ce678cd9
+size 1216062333

text_encoder/config.json

@@ -19,6 +19,6 @@
   "pad_token_id": 1,
   "projection_dim": 768,
   "torch_dtype": "float32",
-  "transformers_version": "4.25.1",
+  "transformers_version": "4.31.0",
   "vocab_size": 49408
 }

text_encoder/pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:57f6e3badaffb5713c93e1f34ac3abf2ee3cd48e60d01714a0a6ed33f3406a5a
-size 492307041
+oid sha256:38a67003cd791d4fc008ae1fd24615b8b168f83cc8e853b746a7ec7bb3d64f42
+size 492306077

tokenizer/tokenizer_config.json

@@ -8,6 +8,7 @@
     "rstrip": false,
     "single_word": false
   },
+  "clean_up_tokenization_spaces": true,
   "do_lower_case": true,
   "eos_token": {
     "__type": "AddedToken",
@@ -19,9 +20,7 @@
   },
   "errors": "replace",
   "model_max_length": 77,
-  "name_or_path": "openai/clip-vit-large-patch14",
   "pad_token": "<|endoftext|>",
-  "special_tokens_map_file": "./special_tokens_map.json",
   "tokenizer_class": "CLIPTokenizer",
   "unk_token": {
     "__type": "AddedToken",