import sys import os import json import zipfile import argparse sys.path.append(os.getcwd()) from os.path import join as opj import numpy as np import torch from PIL import Image, ImageFile from torchvision.transforms import ToPILImage ImageFile.LOAD_TRUNCATED_IMAGES = True import transformers from diffusers import ( StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DPMSolverMultistepScheduler, AutoencoderKL, ) from landmark_generation import generate_annotation from natsort import ns, natsorted class DomainImageGeneration: def __init__( self, src_path, args, model_name, device="cuda:0", clip_skip=2, use_anime_vae=False, save_path='', only_lora=False, model_base_path='./diffusion_model' ): """Initialize DomainImageGeneration class""" self.model_base_path = model_base_path self.device = device self.args = args self.src_path = src_path self.use_model = model_name self.only_lora = only_lora self.out_path_base = os.path.join(save_path, model_name) os.makedirs(self.out_path_base, exist_ok=True) self.diffusion_checkpoint_path = self._get_model_path(model_name) self.pipe = self._load_pipeline(model_name, use_anime_vae, clip_skip) print("All models loaded successfully") def _get_model_path(self, model_name): """Retrieve the model checkpoint path based on the model name""" base_path = self.model_base_path if model_name == "stable-diffusion-2-1-base": return os.path.join(base_path, "stable-diffusion-2-1-base") elif self.only_lora: return os.path.join(base_path, "stable-diffusion-v1-5") else: return os.path.join(base_path, model_name) def _load_controlnet(self, model_name): """Load the ControlNet model""" controlnet_path = os.path.join(self.model_base_path, 'ControlNetMediaPipeFace') if model_name == "stable-diffusion-2-1-base": controlnet_path += "old" return ControlNetModel.from_pretrained( controlnet_path, torch_dtype=torch.float16 ) def _load_pipeline(self, model_name, use_anime_vae, clip_skip): """Load the Stable Diffusion ControlNet Img2Img Pipeline""" controlnet = self._load_controlnet(model_name) if use_anime_vae: print("Using Anime VAE") anime_vae = AutoencoderKL.from_pretrained( "/nas8/liuhongyu/model/kl-f8-anime2", torch_dtype=torch.float16 ) pipeline = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( self.diffusion_checkpoint_path, torch_dtype=torch.float16, safety_checker=None, vae=anime_vae, controlnet=controlnet ).to(self.device) self._load_lora(pipeline, "detail-tweaker-lora/add_detail.safetensors") elif model_name == "stable-diffusion-2-1-base": pipeline = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( self.diffusion_checkpoint_path, torch_dtype=torch.float16, use_safetensors=True, controlnet=controlnet, variant="fp16" ).to(self.device) else: text_encoder = transformers.CLIPTextModel.from_pretrained( self.diffusion_checkpoint_path, subfolder="text_encoder", num_hidden_layers=12 - (clip_skip - 1), torch_dtype=torch.float16 ) pipeline = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( self.diffusion_checkpoint_path, torch_dtype=torch.float16, use_safetensors=True, text_encoder=text_encoder, controlnet=controlnet, variant="fp16" ).to(self.device) self._apply_negative_embedding(pipeline, model_name) pipeline.scheduler = DPMSolverMultistepScheduler.from_config( pipeline.scheduler.config, use_karras_sigmas=True ) print("Target diffusion model loaded") return pipeline def _load_lora(self, pipeline, lora_name): """Load LoRA weights into the model""" lora_path = f"/nas8/liuhongyu/model/{lora_name}" state_dict, network_alphas = pipeline.lora_state_dict(lora_path) pipeline.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=pipeline.unet) def _apply_negative_embedding(self, pipeline, model_name): """Apply negative embedding (textual inversion)""" if model_name not in ["stable-diffusion-xl-base-1.0", "stable-diffusion-2-1-base"]: if self.only_lora: self._load_lora(pipeline, model_name) pipeline.safety_checker = lambda images, clip_input: (images, None) else: pipeline.load_textual_inversion( "/nas8/liuhongyu/lora_model", weight_name="EasyNegativeV2.safetensors", token="EasyNegative" ) def image_generation(self, prompt, strength=0.7, guidance_scale=7.5, num_inference_steps=30): """Generate images using the diffusion model""" out_path = os.path.join(self.out_path_base, prompt.replace(" ", "_")) os.makedirs(out_path, exist_ok=True) src_img_list = natsorted( [f for f in os.listdir(self.src_path) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif'))], alg=ns.PATH ) all_gen_nums = 0 # Counter for generated images for img_name in src_img_list: src_img_pil = Image.open(os.path.join(self.src_path, img_name)) control_image = generate_annotation(src_img_pil, max_faces=1) if control_image is not None: prompt_input = prompt # Apply different generation methods based on the model type if self.use_model in ['stable-diffusion-xl-base-1.0', 'stable-diffusion-2-1-base']: trg_img_pil = self.pipe( prompt=prompt_input, image=src_img_pil, strength=strength, control_image=Image.fromarray(control_image), guidance_scale=guidance_scale, negative_prompt='worst quality, normal quality, low quality, low res, blurry', num_inference_steps=num_inference_steps, controlnet_conditioning_scale=1.5 )['images'][0] else: trg_img_pil = self.pipe( prompt=prompt_input, image=src_img_pil, control_image=Image.fromarray(control_image), strength=strength, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, controlnet_conditioning_scale=1.5, negative_prompt='EasyNegative, worst quality, normal quality, low quality, low res, blurry' )['images'][0] # Save the generated image if valid if np.array(trg_img_pil).max() > 0: trg_img_pil.save(opj(out_path, img_name)) all_gen_nums += 1 def parse_args(): """Parse command-line arguments.""" parser = argparse.ArgumentParser(description="Domain Image Generation") parser.add_argument( "--src_img_path", type=str, default="demo_input", help="Path to the source image directory" ) parser.add_argument("--strength", type=float, default=0.6, help="Strength of the sdeedit") parser.add_argument("--prompt", type=str, default=None, help="Text prompt for image generation") parser.add_argument("--guidance_scale", type=float, default=7.5, help="Guidance scale for Stable Diffusion") parser.add_argument("--sd_model_id", type=str, default="stable-diffusion-2-1-base", help="Stable Diffusion model ID") parser.add_argument("--num_inference_steps", type=int, default=30, help="Number of inference steps") parser.add_argument("--save_base", type=str, default="./output", help="Output directory for generated images") parser.add_argument("--device", type=str, default="cuda:0", help="Device to run inference on (e.g., 'cuda:0')") parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility") parser.add_argument("--use_anime_vae", action="store_true", help="Enable Anime VAE for image generation") parser.add_argument("--model_base_path", type=str, default="./diffusion_model", help="Output directory for generated images") return parser.parse_args() def set_random_seed(seed: int): """Set random seed for reproducibility.""" torch.manual_seed(seed) np.random.seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def main(): """Main function to execute the image generation process.""" args = parse_args() # Set random seed to ensure reproducibility set_random_seed(args.seed) # Check if the source image path exists if not os.path.exists(args.src_img_path): raise FileNotFoundError(f"❌ Source image path does not exist: {args.src_img_path}") # Ensure the output directory exists os.makedirs(args.save_base, exist_ok=True) # Initialize the DomainImageGeneration class and generate images data_generation = DomainImageGeneration( src_path=args.src_img_path, args=args, model_name=args.sd_model_id, save_path=args.save_base, device=args.device, use_anime_vae=args.use_anime_vae, model_base_path = args.model_base_path ) # Start image generation data_generation.image_generation(prompt=args.prompt, strength=args.strength, guidance_scale=args.guidance_scale,num_inference_steps=args.num_inference_steps ) if __name__ == "__main__": main()