Create inference_manager.py

inference_manager.py

@@ -0,0 +1,314 @@
+import spaces
+import os
+import json
+import time
+import copy
+import torch
+from diffusers import AutoPipelineForText2Image, StableDiffusionPipeline,DiffusionPipeline, StableDiffusionXLPipeline, AutoencoderKL, AutoencoderTiny, UNet2DConditionModel
+from huggingface_hub import hf_hub_download, snapshot_download
+from pathlib import Path
+from diffusers import EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DPMSolverSDEScheduler
+from diffusers.models.attention_processor import AttnProcessor2_0
+import os
+from cryptography.hazmat.primitives.asymmetric import rsa, padding
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.asymmetric import utils
+import base64
+import json
+import jwt
+
+#from onediffx import compile_pipe, save_pipe, load_pipe
+
+HF_TOKEN = os.getenv('HF_TOKEN')
+VAR_PUBLIC_KEY = os.getenv('PUBLIC_KEY')
+DATASET_ID = 'nsfwalex/checkpoint_n_lora'
+
+class AuthHelper:
+    def load_public_key_from_file(self):
+        public_key_bytes = VAR_PUBLIC_KEY.encode('utf-8')  # Convert to bytes if it's a string
+        public_key = serialization.load_pem_public_key(
+            public_key_bytes,
+            backend=default_backend()
+        )
+        return public_key
+
+    def __init__(self):
+        self.public_key = self.load_public_key_from_file()
+
+    # check authkey
+    # 1. decode with public key
+    # 2. check timestamp
+    # 3. check current host, referer, ip it should be the same as values in jwt
+
+    def decode_jwt(self, token, algorithms=["RS256"]):
+        """
+        Decode and verify a JWT using a public key.
+        
+        :param public_key: The public key used for verification.
+        :param token: The JWT string to decode.
+        :param algorithms: List of acceptable algorithms (default is ["RS256"]).
+        :return: The decoded JWT payload if verification is successful.
+        :raises: Exception if verification fails.
+        """
+        try:
+            # Decode the JWT
+            decoded_payload = jwt.decode(
+                token,
+                self.public_key,
+                algorithms=algorithms,
+                options={"verify_signature": True}  # Explicitly enable signature verification
+            )
+            return decoded_payload
+        except Exception as e:
+            print("Invalid token:", e)
+            raise
+
+    def check_auth(self, session, token):
+        params = session.get("params") or {}
+        if params.get("_skip_token_passkey", "") == "nsfwaisio_125687":
+            return True
+        sip = session.get("client_ip", "")
+        shost = session.get("host", "")
+        sreferer = session.get("refer")
+        print(sip, shost, sreferer)
+        jwt_data = self.decode_jwt(token)
+        tip = jwt_data.get("ip", "")
+        thost = jwt_data.get("host", "")
+        treferer = jwt_data.get("referer", "")
+        print(sip, tip, shost, thost, sreferer, treferer)
+        if not tip or not thost or not treferer:
+            raise Exception("invalid token")
+        if sip == tip and shost == thost and sreferer == treferer:
+            return True
+        raise Exception("wrong token")
+
+class InferenceManager:
+    def __init__(self, model_version="xl", config_path="config.json", lora_options_path="loras.json"):
+        self.model_version = model_version
+        self.lora_load_options = self.load_json(lora_options_path)  # Load LoRA load options
+        self.lora_models = self.load_index_file("index.json")  # Load index.json
+        self.preloaded_loras = []  # Array to store preloaded LoRAs with name and weights
+        self.base_model_pipeline = self.load_base_model(config_path)  # Load the base model
+        self.preload_loras()  # Preload LoRAs based on options
+
+    def load_json(self, filepath):
+        """Load JSON file into a dictionary."""
+        if os.path.exists(filepath):
+            with open(filepath, "r", encoding="utf-8") as f:
+                return json.load(f)
+        return {}
+
+    def load_index_file(self, index_file):
+        """Download index.json from Hugging Face and return the file path."""
+        index_path = download_from_hf(index_file)
+        if index_path:
+            with open(index_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        return {}
+
+    @spaces.GPU(duration=40)
+    def compile_onediff(self):
+        self.base_model_pipeline.to("cuda")
+        pipe = self.base_model_pipeline
+        # load the compiled pipe
+        load_pipe(pipe, dir="cached_pipe")
+        print("Start oneflow compiling...")
+        start_compile = time.time()
+        pipe = compile_pipe(pipe)
+        # run once to trigger compilation
+        image = pipe(
+            prompt="street style, detailed, raw photo, woman, face, shot on CineStill 800T",
+            height=512,
+            width=512,
+            num_inference_steps=10,
+            output_type="pil",
+        ).images
+        image[0].save(f"test_image.png")
+        compile_time = time.time() - start_compile
+        #self.base_model_pipeline.to("cpu")
+        # save the compiled pipe
+        save_pipe(pipe, dir="cached_pipe")
+        self.base_model_pipeline = pipe
+        print(f"OneDiff compile in {compile_time}s")
+    
+    def load_base_model(self, config_path):
+        """Load base model and return the pipeline."""
+        start = time.time()
+        with open(config_path, "r", encoding="utf-8") as f:
+            cfg = json.load(f)
+
+        model_version = cfg.get("model_version", self.model_version)
+        ckpt_dir = snapshot_download(repo_id=cfg["model_id"], local_files_only=False)
+
+        if model_version == "1.5":
+            vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.bfloat16)
+            pipe = StableDiffusionPipeline.from_pretrained(ckpt_dir, vae=vae, torch_dtype=torch.bfloat16, use_safetensors=True)
+        else:
+            
+            #vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.bfloat16)
+            vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.bfloat16)
+            
+            pipe = DiffusionPipeline.from_pretrained(
+                ckpt_dir,
+                vae=vae,
+                #unet=unet,
+                torch_dtype=torch.bfloat16,
+                use_safetensors=True,
+                variant="fp16",
+                custom_pipeline   = "lpw_stable_diffusion_xl",
+            )
+            
+            clip_skip = cfg.get("clip_skip", 1)
+            # Adjust clip skip for XL (assumed not relevant for SD 1.5)
+            pipe.text_encoder.config.num_hidden_layers -= (clip_skip - 1)
+        
+        load_time = round(time.time() - start, 2)
+        print(f"Base model loaded in {load_time}s")
+        return pipe
+
+    def preload_loras(self):
+        """Preload all LoRAs marked as 'preload=True' and store for later use."""
+        for lora_name, lora_info in self.lora_load_options.items():
+            try:
+                start = time.time()
+
+                # Find the corresponding LoRA in index.json
+                lora_index_info = next((l for l in self.lora_models['lora'] if l['name'] == lora_name), None)
+                if not lora_index_info:
+                    raise ValueError(f"LoRA {lora_name} not found in index.json.")
+
+                # Check if the LoRA base model matches the current model version
+                if self.model_version not in lora_info['base_model'] or not lora_info.get('preload', False):
+                    print(f"Skipping {lora_name} as it's not compatible with the current model version.")
+                    continue
+
+                # Load LoRA weights from the specified path
+                weight_path = download_from_hf(lora_index_info['path'], local_dir=None)
+                if not weight_path:
+                    raise ValueError(f"Failed to download LoRA weights for {lora_name}")
+                load_time = round(time.time() - start, 2)
+                print(f"Downloaded {lora_name} in {load_time}s")
+                self.base_model_pipeline.load_lora_weights(
+                    weight_path,
+                    weight_name=lora_index_info["path"],
+                    adapter_name=lora_name
+                )
+
+                # Store the preloaded LoRA name and weight for merging later
+                if lora_info.get("preload", False):
+                    self.preloaded_loras.append({
+                        "name": lora_name,
+                        "weight": lora_info.get("weight", 1.0)
+                    })
+                    load_time = round(time.time() - start, 2)
+                    print(f"Preloaded LoRA {lora_name} with weight {lora_info.get('weight', 1.0)} in {load_time}s.")
+            except Exception as e:
+                print(f"Lora {lora_name} not loaded, skipping... {e}")
+
+    def build_pipeline_with_lora(self, lora_list, sampler="DPM2 a", new_pipeline=False):
+        """Build the pipeline with specific LoRAs, loading any that are not preloaded."""
+        # Deep copy the base pipeline
+        start = time.time()
+        if new_pipeline:
+            temp_pipeline = copy.deepcopy(self.base_model_pipeline)
+        else:
+            temp_pipeline = self.base_model_pipeline
+        copy_time = round(time.time() - start, 2)
+        print(f"pipeline copied in {copy_time}s")
+        # Track LoRAs to be loaded dynamically
+        dynamic_loras = []
+
+        # Check if any LoRAs in lora_list need to be loaded dynamically
+        for lora_name in lora_list:
+            if not any(l['name'] == lora_name for l in self.preloaded_loras):
+                lora_info = next((l for l in self.lora_models['lora'] if l['name'] == lora_name), None)
+                if lora_info and self.model_version in lora_info["attr"].get("base_model", []):
+                    dynamic_loras.append({
+                        "name": lora_name,
+                        "filename": lora_info["path"],
+                        "scale": 1.0  # Assuming default weight as 1.0 for dynamic LoRAs
+                    })
+
+        # Fuse preloaded and dynamic LoRAs
+        all_loras = [{"name": x["name"], "scale": x["weight"], "preloaded": True} for x in self.preloaded_loras] + dynamic_loras
+        set_lora_weights(temp_pipeline, all_loras,False)
+
+        build_time = round(time.time() - start, 2)
+        print(f"Pipeline built with LoRAs in {build_time}s.")
+
+        # Define samplers
+        samplers = {
+            "Euler a": EulerAncestralDiscreteScheduler.from_config(temp_pipeline.scheduler.config),
+            "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(temp_pipeline.scheduler.config, use_karras_sigmas=True),
+            "DPM2 a": DPMSolverMultistepScheduler.from_config(temp_pipeline.scheduler.config)
+        }
+        
+        # Set the scheduler based on the selected sampler
+        temp_pipeline.scheduler = samplers[sampler]
+
+        # Move the final pipeline to the GPU
+        temp_pipeline
+        return temp_pipeline
+
+    def release(self, temp_pipeline):
+        """Release the deepcopied pipeline to recycle memory."""
+        del temp_pipeline
+        torch.cuda.empty_cache()
+        print("Memory released and cache cleared.")
+
+
+# Hugging Face file download function - returns only file path
+def download_from_hf(filename, local_dir=None):
+    try:
+        file_path = hf_hub_download(
+            filename=filename,
+            repo_id=DATASET_ID,
+            repo_type="dataset",
+            revision="main",
+            local_dir=local_dir,
+            local_files_only=False,  # Attempt to load from cache if available
+        )
+        return file_path  # Return file path only
+    except Exception as e:
+        print(f"Failed to load {filename} from Hugging Face: {str(e)}")
+        return None
+
+
+# Function to load and fuse LoRAs
+def set_lora_weights(pipe, lorajson: list[dict], fuse=False):
+    try:
+        if not lorajson or not isinstance(lorajson, list):
+            return
+
+        a_list = []
+        w_list = []
+        for d in lorajson:
+            if not d or not isinstance(d, dict) or not d["name"] or d["name"] == "None":
+                continue
+
+            k = d["name"]
+            if not d.get("preloaded", False):
+                start = time.time()
+                weight_path = download_from_hf(d['filename'], local_dir=None)
+                if weight_path:
+                    pipe.load_lora_weights(weight_path, weight_name=d['filename'], adapter_name=k)
+
+                load_time = round(time.time() - start, 2)
+                print(f"LoRA {k} loaded in {load_time}s.")
+
+            a_list.append(k)
+            w_list.append(d["scale"])
+
+        if not a_list:
+            return
+
+        start = time.time()
+        pipe.set_adapters(a_list, adapter_weights=w_list)
+        if fuse:
+            pipe.fuse_lora(adapter_names=a_list, lora_scale=1.0)
+        fuse_time = round(time.time() - start, 2)
+        print(f"LoRAs fused in {fuse_time}s.")
+    except Exception as e:
+        print(f"External LoRA Error: {e}")
+        raise Exception(f"External LoRA Error: {e}") from e