Upload web-ui.py

web-ui.py

@@ -0,0 +1,172 @@
+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
+from ip_adapter.ip_adapter_faceid import IPAdapterFaceIDPlus
+import argparse
+import random
+from insightface.utils import face_align
+
+# 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")
+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",
+	r"C:\Users\King\Downloads\New folder\3D Animation Diffusion"
+]
+
+# Cache for loaded models
+model_cache = {}
+max_cache_size = args.cache_limit
+
+# Function to load and cache model
+def load_model(model_name):
+    if model_name in model_cache:
+        return model_cache[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"
+    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)
+
+    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):
+    saved_images = []
+
+    # Load and prepare the model
+    ip_model = load_model(model_name)
+
+    # Convert input image to the format expected by the model
+    input_image = input_image.convert("RGB")
+    input_image = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+    app = FaceAnalysis(
+        name="buffalo_l", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
+    )
+    app.prepare(ctx_id=0, det_size=(640, 640))
+    faces = app.get(input_image)
+    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, landmark=faces[0].kps, image_size=224)
+
+    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,
+            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):
+            image_path = os.path.join(outputs_dir, f"generated_{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")
+        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():            
+            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])
+
+    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)
+
+    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],
+        outputs=[output_gallery, output_text, display_seed],
+    )
+
+demo.launch(share=args.share, inbrowser=True)