web-ui.py

import gradio as gr
import os
import cv2
import numpy as np
import torch
from PIL import Image
from insightface.app import FaceAnalysis
from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL, StableDiffusionXLPipeline
from ip_adapter.ip_adapter_faceid import IPAdapterFaceIDPlus
import argparse
import random
from insightface.utils import face_align
from pyngrok import ngrok
import threading
import time
from ip_adapter.ip_adapter_faceid import IPAdapterFaceIDPlusXL
import hashlib
from datetime import datetime

# Argument parser for command line options
parser = argparse.ArgumentParser()
parser.add_argument("--share", action="store_true", help="Enable Gradio share option")
parser.add_argument("--num_images", type=int, default=1, help="Number of images to generate")
parser.add_argument("--cache_limit", type=int, default=1, help="Limit for model cache")
parser.add_argument("--ngrok_token", type=str, default=None, help="ngrok authtoken for tunneling")

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",    
    "SG161222/Realistic_Vision_V2.0",
    "SG161222/Realistic_Vision_V4.0_noVAE",
    "SG161222/Realistic_Vision_V5.1_noVAE",
    "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

embeddings_cache = {}

def get_image_hash(image):
    image_bytes = image.tobytes()
    return hashlib.sha256(image_bytes).hexdigest()

def convert_model(checkpoint_path, output_path, isSDXL):
    try:
        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, isSDXL):
    if model_name in model_cache:
        return model_cache[model_name]
    print(f"loading model {model_name}")
    # Limit cache size
    if len(model_cache) >= max_cache_size:
        model_cache.pop(next(iter(model_cache)))

    device = "cuda"
    noise_scheduler = DDIMScheduler(
        num_train_timesteps=1000,
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule="scaled_linear",
        clip_sample=False,
        set_alpha_to_one=False,
        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)

    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:
        image_encoder_path = "h94/IP-Adapter/models/image_encoder"
        ip_ckpt = "adapters/ip-adapter-faceid-plusv2_sdxl.bin"
        ip_model = IPAdapterFaceIDPlusXL(pipe,image_encoder_path, 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, isSDXL,cfg):
    saved_images = []
    if custom_model_path:
        model_name = custom_model_path
    # Load and prepare the model
    ip_model = load_model(model_name, isSDXL)

    # Convert input image to the format expected by the model and calculate its hash
    input_image = input_image.convert("RGB")
    input_image_cv2 = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
    image_hash = get_image_hash(input_image)

    # Check if embeddings are cached
    if image_hash in embeddings_cache:
        faceid_embeds, face_image = embeddings_cache[image_hash]
    else:
        app = FaceAnalysis(
            name="buffalo_l", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
        )
        app.prepare(ctx_id=0, det_size=(640, 640))
        faces = app.get(input_image_cv2)
        if not faces:
            raise ValueError("No faces found in the image.")

        faceid_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
        face_image = face_align.norm_crop(input_image_cv2, landmark=faces[0].kps, image_size=224)
        # Cache the embeddings
        embeddings_cache[image_hash] = (faceid_embeds, face_image)

    for image_index in range(num_images):
        if randomize_seed or image_index > 0:
            seed = random.randint(0, 2**32 - 1)

        # Generate the image with the new parameters
        generated_images = ip_model.generate(
            prompt=positive_prompt,
            negative_prompt=negative_prompt,
            faceid_embeds=faceid_embeds,
            face_image=face_image,
            num_samples=batch_size,
            shortcut=enable_shortcut,
            s_scale=s_scale,
            width=width,
            height=height,
			guidance_scale=cfg,
            num_inference_steps=num_inference_steps,
            seed=seed,
        )

        # Save and prepare the generated images for display
        outputs_dir = "outputs"
        if not os.path.exists(outputs_dir):
            os.makedirs(outputs_dir)
        for i, img in enumerate(generated_images, start=1):
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            image_path = os.path.join(outputs_dir, f"{timestamp}_image_{len(os.listdir(outputs_dir)) + i}.png")
            img.save(image_path)
            saved_images.append(image_path)

    return saved_images, f"Saved images: {', '.join(saved_images)}", seed

# Gradio interface, using the static list of models
with gr.Blocks() as demo:
    gr.Markdown("Developed by SECourses - only distributed on https://www.patreon.com/posts/95759342")
    with gr.Row():
        input_image = gr.Image(type="pil")
        generate_btn = gr.Button("Generate")
        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")
            randomize_seed = gr.Checkbox(value=True, label="Randomize Seed")
        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():  
            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():
        positive_prompt = gr.Textbox(label="Positive Prompt")
        negative_prompt = gr.Textbox(label="Negative Prompt")    
    with gr.Row():            
        model_selector = gr.Dropdown(label="Select Model", choices=static_model_names, value=static_model_names[0])
        custom_model_path = gr.Textbox(label="Custom Model Path (Optional)")

    with gr.Column():
        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, isSDXL,cfg],
        outputs=[output_gallery, output_text, display_seed]
    )
	
    convert_btn.click(
        convert_model,
        inputs=[checkpoint_path_input, output_path_input, isSDXL],
        outputs=[gr.Text(label="Conversion Status")],
    )	

# Function to start ngrok for tunneling
def start_ngrok():
    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 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)