Upload web-ui.py

web-ui.py

@@ -5,7 +5,7 @@ import numpy as np
 import torch
 from PIL import Image
 from insightface.app import FaceAnalysis
-from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
+from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL, StableDiffusionXLPipeline
 from ip_adapter.ip_adapter_faceid import IPAdapterFaceIDPlus
 import argparse
 import random
@@ -13,6 +13,7 @@ from insightface.utils import face_align
 from pyngrok import ngrok
 import threading
 import time
+from ip_adapter.ip_adapter_faceid import IPAdapterFaceIDXL
 
 # Argument parser for command line options
 parser = argparse.ArgumentParser()
@@ -26,38 +27,45 @@ args = parser.parse_args()
 # Add new model names here
 static_model_names = [
     "SG161222/Realistic_Vision_V6.0_B1_noVAE",
-	"stablediffusionapi/rev-animated-v122-eol",
-	"Lykon/DreamShaper",
-	"stablediffusionapi/toonyou",
-	"stablediffusionapi/real-cartoon-3d",
-	"KBlueLeaf/kohaku-v2.1",
-	"nitrosocke/Ghibli-Diffusion",
-	"Linaqruf/anything-v3.0",
-	"jinaai/flat-2d-animerge",
-	"stablediffusionapi/realcartoon3d",
-	"stablediffusionapi/disney-pixar-cartoon",
-	"stablediffusionapi/pastel-mix-stylized-anime",
-	"stablediffusionapi/anything-v5",	
+    "stablediffusionapi/rev-animated-v122-eol",
+    "Lykon/DreamShaper",
+    "stablediffusionapi/toonyou",
+    "stablediffusionapi/real-cartoon-3d",
+    "KBlueLeaf/kohaku-v2.1",
+    "nitrosocke/Ghibli-Diffusion",
+    "Linaqruf/anything-v3.0",
+    "jinaai/flat-2d-animerge",
+    "stablediffusionapi/realcartoon3d",
+    "stablediffusionapi/disney-pixar-cartoon",
+    "stablediffusionapi/pastel-mix-stylized-anime",
+    "stablediffusionapi/anything-v5",    
     "SG161222/Realistic_Vision_V2.0",
     "SG161222/Realistic_Vision_V4.0_noVAE",
     "SG161222/Realistic_Vision_V5.1_noVAE",
-	#r"G:\model\model_diffusers"
+    "stablediffusionapi/anime-illust-diffusion-xl",
+	"stabilityai/stable-diffusion-xl-base-1.0",
+    #r"G:\model\model_diffusers"
 ]
 
 # Cache for loaded models
 model_cache = {}
 max_cache_size = args.cache_limit
 
-def convert_model(checkpoint_path, output_path):
+def convert_model(checkpoint_path, output_path, isSDXL):
     try:
-        pipe = StableDiffusionPipeline.from_single_file(checkpoint_path)
-        pipe.save_pretrained(output_path)
+        if isSDXL:
+            pipe = StableDiffusionXLPipeline.from_single_file(checkpoint_path)
+            pipe.save_pretrained(output_path)
+        else:
+            pipe = StableDiffusionPipeline.from_single_file(checkpoint_path)
+            pipe.save_pretrained(output_path)
         return f"Model converted and saved to {output_path}"
     except Exception as e:
         return f"Error: {str(e)}"
 
+
 # Function to load and cache model
-def load_model(model_name):
+def load_model(model_name, isSDXL):
     if model_name in model_cache:
         return model_cache[model_name]
     print(f"loading model {model_name}")
@@ -76,32 +84,47 @@ def load_model(model_name):
         steps_offset=1,
     )
     vae_model_path = "stabilityai/sd-vae-ft-mse"
+    if isSDXL:
+        vae_model_path = "stabilityai/sdxl-vae"
     vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
 
-    # Load model based on the selected model name
-    pipe = StableDiffusionPipeline.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16,
-        scheduler=noise_scheduler,
-        vae=vae,
-        feature_extractor=None,
-        safety_checker=None
-    ).to(device)
-
-    image_encoder_path = "h94/IP-Adapter/models/image_encoder"
-    ip_ckpt = "adapters/ip-adapter-faceid-plusv2_sd15.bin"
-    ip_model = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_ckpt, device)
+    if isSDXL:
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            vae=vae,
+            scheduler=noise_scheduler,
+            add_watermarker=False,
+        ).to(device)
+    else:
+        # Load model based on the selected model name
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            scheduler=noise_scheduler,
+            vae=vae,
+            feature_extractor=None,
+            safety_checker=None
+        ).to(device)
+
+    if isSDXL:
+        ip_ckpt = "adapters/ip-adapter-faceid_sdxl.bin"
+        ip_model = IPAdapterFaceIDXL(pipe, ip_ckpt, device)
+    else:
+        image_encoder_path = "h94/IP-Adapter/models/image_encoder"
+        ip_ckpt = "adapters/ip-adapter-faceid-plusv2_sd15.bin"
+        ip_model = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_ckpt, device)
 
     model_cache[model_name] = ip_model
     return ip_model
 
 # Function to process image and generate output
-def generate_image(input_image, positive_prompt, negative_prompt, width, height, model_name, num_inference_steps, seed, randomize_seed, num_images, batch_size, enable_shortcut, s_scale, custom_model_path):
+def generate_image(input_image, positive_prompt, negative_prompt, width, height, model_name, num_inference_steps, seed, randomize_seed, num_images, batch_size, enable_shortcut, s_scale, custom_model_path, isSDXL,cfg):
     saved_images = []
     if custom_model_path:
         model_name = custom_model_path
     # Load and prepare the model
-    ip_model = load_model(model_name)
+    ip_model = load_model(model_name, isSDXL)
 
     # Convert input image to the format expected by the model
     input_image = input_image.convert("RGB")
@@ -132,6 +155,7 @@ def generate_image(input_image, positive_prompt, negative_prompt, width, height,
             s_scale=s_scale,
             width=width,
             height=height,
+			guidance_scale=cfg,
             num_inference_steps=num_inference_steps,
             seed=seed,
         )
@@ -156,6 +180,7 @@ with gr.Blocks() as demo:
         with gr.Row():
             width = gr.Number(value=512, label="Width")
             height = gr.Number(value=768, label="Height")
+            cfg = gr.Number(value=7.5, label="CFG")
         with gr.Row():        
             num_inference_steps = gr.Number(value=30, label="Number of Inference Steps", step=1, minimum=10, maximum=100)
             seed = gr.Number(value=2023, label="Seed")
@@ -163,7 +188,8 @@ with gr.Blocks() as demo:
         with gr.Row():                    
             num_images = gr.Number(value=args.num_images, label="Number of Images to Generate", step=1, minimum=1)
             batch_size = gr.Number(value=1, label="Batch Size", step=1)
-        with gr.Row():            
+        with gr.Row():  
+            isSDXL = gr.Checkbox(value=False, label="Activate SDXL")
             enable_shortcut = gr.Checkbox(value=True, label="Enable Shortcut")
             s_scale = gr.Number(value=1.0, label="Scale Factor (s_scale)", step=0.1, minimum=0.5, maximum=4.0)
     with gr.Row():
@@ -177,39 +203,37 @@ with gr.Blocks() as demo:
         output_gallery = gr.Gallery(label="Generated Images")
         output_text = gr.Textbox(label="Output Info")
         display_seed = gr.Textbox(label="Used Seed", interactive=False)
-		
+        
     with gr.Row():
         checkpoint_path_input = gr.Textbox(label="Enter Checkpoint File Path .e.g G:\model\model.safetensors", )
         output_path_input = gr.Textbox(label="Enter Output Folder Path, e.g. G:\model\model_diffusers")
         convert_btn = gr.Button("Convert Model")
 
-
     generate_btn.click(
         generate_image,
-        inputs=[input_image, positive_prompt, negative_prompt, width, height, model_selector, num_inference_steps, seed, randomize_seed, num_images, batch_size, enable_shortcut, s_scale, custom_model_path],
+        inputs=[input_image, positive_prompt, negative_prompt, width, height, model_selector, num_inference_steps, seed, randomize_seed, num_images, batch_size, enable_shortcut, s_scale, custom_model_path, isSDXL,cfg],
         outputs=[output_gallery, output_text, display_seed]
     )
 	
     convert_btn.click(
         convert_model,
-        inputs=[checkpoint_path_input, output_path_input],
+        inputs=[checkpoint_path_input, output_path_input, isSDXL],
         outputs=[gr.Text(label="Conversion Status")],
     )	
-	
-#sadly doesnt work
+
+# Function to start ngrok for tunneling
 def start_ngrok():
-    print("1")
-    time.sleep(10)  # Delay for 10 seconds to ensure Gradio starts first
-    print("2")
+    print("Starting ngrok...")
+    time.sleep(10)  # Delay to ensure Gradio starts first
     ngrok.set_auth_token(args.ngrok_token)
     public_url = ngrok.connect(port=7860)  # Adjust to your Gradio app's port
     print(f"ngrok tunnel started at {public_url}")
 
 if __name__ == "__main__":
-    #if args.ngrok_token:
-        # Start ngrok in a daemon thread with a delay
-    #    ngrok_thread = threading.Thread(target=start_ngrok, daemon=False)
-    #    ngrok_thread.start()
+    if args.ngrok_token:
+        # Start ngrok in a separate thread with a delay
+        ngrok_thread = threading.Thread(target=start_ngrok, daemon=True)
+        ngrok_thread.start()
 
     # Launch the Gradio app
     demo.launch(share=args.share, inbrowser=True)