File size: 7,461 Bytes
a6f8521 3857a74 49b3a73 bfa3982 a6f8521 3857a74 a6f8521 3857a74 bf5edb4 bfa3982 bf5edb4 49b3a73 bf5edb4 bfa3982 bf5edb4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | 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
from pyngrok import ngrok
import threading
import time
# 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",
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],
)
def start_ngrok():
time.sleep(10) # Delay for 10 seconds 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=True)
ngrok_thread.start()
# Launch the Gradio app
demo.launch(share=args.share, inbrowser=True)
|