MonsterMMORPG
/

Stable-Diffusion

Model card Files Files and versions Community

Stable-Diffusion / web-ui.py

MonsterMMORPG

Upload web-ui.py

006099c verified 4 months ago

raw history blame contribute delete

No virus

10.3 kB

	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)