Update inference_manager.py

inference_manager.py

@@ -27,12 +27,29 @@ import re
 import gradio as gr
 import uuid
 from PIL import Image
-MAX_SEED = 12211231#np.iinfo(np.int32).max
+MAX_SEED = np.iinfo(np.int32).max
 #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'
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "beta_start": 0.00085,
+    "beta_end": 0.012,
+    "beta_schedule": "scaled_linear",
+    "set_alpha_to_one": False,
+    "steps_offset": 1,
+    "prediction_type": "epsilon",
+}
+samplers = {
+    "Euler a": EulerAncestralDiscreteScheduler.from_config(scheduler_config),
+    "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(scheduler_config, use_karras_sigmas=True),
+    "DPM2 a": DPMSolverMultistepScheduler.from_config(scheduler_config),
+    "DPM++ SDE": DPMSolverSDEScheduler.from_config(scheduler_config),
+    "DPM++ 2M SDE": DPMSolverSDEScheduler.from_config(scheduler_config, use_2m=True),
+    "DPM++ 2S a": DPMSolverMultistepScheduler.from_config(scheduler_config, use_2s=True)
+}
 
 class AuthHelper:
     def load_public_key_from_file(self):
@@ -179,7 +196,7 @@ class InferenceManager:
             pipe = StableDiffusionPipeline.from_pretrained(ckpt_dir, vae=vae, torch_dtype=torch.bfloat16, use_safetensors=True)
         else:
             use_vae = cfg.get("vae", "")
-            if not use_vae:
+            if not use_vae or True:#!TEST! default vae for test
                 vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.bfloat16)
             elif use_vae == "tae":
                 vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.bfloat16)
@@ -193,7 +210,7 @@ class InferenceManager:
                 torch_dtype=torch.bfloat16,
                 use_safetensors=True,
                 #variant="fp16",
-                #custom_pipeline   = "lpw_stable_diffusion_xl",
+                custom_pipeline   = "lpw_stable_diffusion_xl",
             )
             #pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
             clip_skip = cfg.get("clip_skip", 1)
@@ -208,7 +225,7 @@ class InferenceManager:
                 ip_ckpt = self.ext_model_pathes.get("ip-adapter-faceid-sdxl", "")
                 if ip_ckpt:
                     print(f"loading ip adapter model...")
-                    self.ip_adapter_faceid_pipeline = ipown.IPAdapterFaceIDXL(pipe, ip_ckpt, 'cuda')
+                    self.ip_adapter_faceid_pipeline = ipown.IPAdapterFaceIDXL(pipe, ip_ckpt, 'cuda', torch_dtype=torch.bfloat16)
                 else:
                     print("ip-adapter-faceid-sdxl not found, skip")            
 
@@ -583,8 +600,8 @@ class ModelManager:
             generator=generator,
             num_images_per_prompt=1,
             output_type="pil",
-            callback_on_step_end=callback_dynamic_cfg,
-            callback_on_step_end_tensor_inputs=['prompt_embeds', 'add_text_embeds', 'add_time_ids'],
+            #callback_on_step_end=callback_dynamic_cfg,
+            #callback_on_step_end_tensor_inputs=['prompt_embeds', 'add_text_embeds', 'add_time_ids'],
         ).images
         cost = round(time.time() - start, 2)
         print(f"inference done in {cost}s")