import gradio as gr import numpy as np import torch from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline, StableDiffusionControlNetImg2ImgPipeline from peft import PeftModel, LoraConfig import os from PIL import Image MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_default = "stable-diffusion-v1-5/stable-diffusion-v1-5" torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 # Инициализация ControlNet controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-openpose", torch_dtype=torch_dtype) def get_lora_sd_pipeline( lora_dir='./lora_man_animestyle', base_model_name_or_path=None, dtype=torch.float16, adapter_name="default" ): unet_sub_dir = os.path.join(lora_dir, "unet") text_encoder_sub_dir = os.path.join(lora_dir, "text_encoder") if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None: config = LoraConfig.from_pretrained(text_encoder_sub_dir) base_model_name_or_path = config.base_model_name_or_path if base_model_name_or_path is None: raise ValueError("Укажите название базовой модели или путь к ней") pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype) before_params = pipe.unet.parameters() pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name) pipe.unet.set_adapter(adapter_name) after_params = pipe.unet.parameters() if os.path.exists(text_encoder_sub_dir): pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name) if dtype in (torch.float16, torch.bfloat16): pipe.unet.half() pipe.text_encoder.half() return pipe def long_prompt_encoder(prompt, tokenizer, text_encoder, max_length=77): tokens = tokenizer(prompt, truncation=False, return_tensors="pt")["input_ids"] part_s = [tokens[:, i:i + max_length] for i in range(0, tokens.shape[1], max_length)] with torch.no_grad(): embeds = [text_encoder(part.to(text_encoder.device))[0] for part in part_s] return torch.cat(embeds, dim=1) def align_embeddings(prompt_embeds, negative_prompt_embeds): max_length = max(prompt_embeds.shape[1], negative_prompt_embeds.shape[1]) return torch.nn.functional.pad(prompt_embeds, (0, 0, 0, max_length - prompt_embeds.shape[1])), \ torch.nn.functional.pad(negative_prompt_embeds, (0, 0, 0, max_length - negative_prompt_embeds.shape[1])) def preprocess_image(image, target_width, target_height): # Преобразует изображение в формат, подходящий для модели. if isinstance(image, np.ndarray): image = Image.fromarray(image) image = image.resize((target_width, target_height), Image.LANCZOS) image = np.array(image).astype(np.float32) / 255.0 # Нормализация [0, 1] image = image[None].transpose(0, 3, 1, 2) # Преобразуем в (batch, channels, height, width) image = torch.from_numpy(image).to(device) return image pipe_default = get_lora_sd_pipeline(lora_dir='./lora_man_animestyle', base_model_name_or_path=model_default, dtype=torch_dtype).to(device) #pipe_controlnet = StableDiffusionControlNetPipeline.from_pretrained( pipe_controlnet = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( model_default, controlnet=controlnet, torch_dtype=torch_dtype ).to(device) def infer( prompt, negative_prompt, width=512, height=512, num_inference_steps=20, model='stable-diffusion-v1-5/stable-diffusion-v1-5', seed=4, guidance_scale=7.5, lora_scale=0.5, strength_cn=0.5, # Коэфф. зашумления ControlNet use_control_net=False, # Параметр для включения ControlNet control_strength=0.5, # Сила влияния ControlNet cn_source_image=None, # Исходное изображение ControlNet control_image=None, # Контрольное изображение ControlNet strength_ip=0.5, # Коэфф. зашумления IP_adapter use_ip_adapter=False, # Параметр для включения IP_adapter ip_adapter_strength=0.5,# Сила влияния IP_adapter ip_source_image=None, # Исходное изображение IP_adapter ip_adapter_image=None, # Контрольное изображение IP_adapter progress=gr.Progress(track_tqdm=True) ): generator = torch.Generator(device).manual_seed(seed) # Генерация с IP_adapter if use_ip_adapter and ip_source_image is not None and ip_adapter_image is not None: # pipe_controlnet = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( # model_default, # controlnet=controlnet, # torch_dtype=torch_dtype # ).to(device) # Преобразуем изображения ip_source_image = preprocess_image(ip_source_image, width, height) ip_adapter_image = preprocess_image(ip_adapter_image, width, height) # Создаём пайплайн IP_adapter с LoRA, если он ещё не создан ??????????????????????????????????????????????????????????????? if not hasattr(pipe_controlnet, 'lora_loaded') or not pipe_controlnet.lora_loaded: # Загружаем LoRA для UNet pipe_controlnet.unet = PeftModel.from_pretrained( pipe_controlnet.unet, './lora_man_animestyle/unet', adapter_name="default" ) pipe_controlnet.unet.set_adapter("default") # Загружаем LoRA для Text Encoder, если она существует text_encoder_lora_path = './lora_man_animestyle/text_encoder' if os.path.exists(text_encoder_lora_path): pipe_controlnet.text_encoder = PeftModel.from_pretrained( pipe_controlnet.text_encoder, text_encoder_lora_path, adapter_name="default" ) pipe_controlnet.text_encoder.set_adapter("default") # Объединяем LoRA с основной моделью pipe_controlnet.fuse_lora(lora_scale=lora_scale) pipe_controlnet.lora_loaded = True # Помечаем, что LoRA загружена # Убедимся, что ip_adapter_strength имеет тип float ip_adapter_strength = float(ip_adapter_strength) #strength_ip = float(strength_ip) # Используем IP_adapter с LoRA ???????????????????????????????????????????????????????????????????????? pipe = pipe_controlnet prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder) negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder) prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds) image = pipe_controlnet( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, image=ip_source_image, control_image=ip_adapter_image, strength=strength_ip, # Коэфф. зашумления, чем больше, тем больше меняется результирующее изображение относитенльно исходного width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, controlnet_conditioning_scale=ip_adapter_strength, # ??????????????????????????????????????????????????????????????? generator=generator ).images[0] else: # Генерация с ControlNet if use_control_net and control_image is not None and cn_source_image is not None: # pipe_controlnet = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( # model_default, # controlnet=controlnet, # torch_dtype=torch_dtype # ).to(device) # Преобразуем изображения cn_source_image = preprocess_image(cn_source_image, width, height) control_image = preprocess_image(control_image, width, height) # Создаём пайплайн ControlNet с LoRA, если он ещё не создан if not hasattr(pipe_controlnet, 'lora_loaded') or not pipe_controlnet.lora_loaded: # Загружаем LoRA для UNet pipe_controlnet.unet = PeftModel.from_pretrained( pipe_controlnet.unet, './lora_man_animestyle/unet', adapter_name="default" ) pipe_controlnet.unet.set_adapter("default") # Загружаем LoRA для Text Encoder, если она существует text_encoder_lora_path = './lora_man_animestyle/text_encoder' if os.path.exists(text_encoder_lora_path): pipe_controlnet.text_encoder = PeftModel.from_pretrained( pipe_controlnet.text_encoder, text_encoder_lora_path, adapter_name="default" ) pipe_controlnet.text_encoder.set_adapter("default") # Объединяем LoRA с основной моделью pipe_controlnet.fuse_lora(lora_scale=lora_scale) pipe_controlnet.lora_loaded = True # Помечаем, что LoRA загружена # Убедимся, что control_strength имеет тип float control_strength = float(control_strength) #strength_sn = float(strength_sn) # Используем ControlNet с LoRA pipe = pipe_controlnet prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder) negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder) prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds) image = pipe_controlnet( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, image=cn_source_image, control_image=control_image, strength=strength_cn, # Коэфф. зашумления, чем больше, тем больше меняется результирующее изображение относитенльно исходного width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, controlnet_conditioning_scale=control_strength, generator=generator ).images[0] else: # Генерация без ControlNet и IP_adapter if model != model_default: pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device) prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder) negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder) prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds) else: pipe = pipe_default prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder) negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder) prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds) pipe.fuse_lora(lora_scale=lora_scale) params = { 'prompt_embeds': prompt_embeds, 'negative_prompt_embeds': negative_prompt_embeds, 'guidance_scale': guidance_scale, 'num_inference_steps': num_inference_steps, 'width': width, 'height': height, 'generator': generator, } image = pipe(**params).images[0] return image examples = [ "A young man in anime style. The image is characterized by high definition and resolution. Handsome, thoughtful man, attentive eyes. The man is depicted in the foreground, close-up or in the middle. High-quality images of the face, eyes, nose, lips, hands and clothes. The background and background are blurred and indistinct. The play of light and shadow is visible on the face and clothes.", "A man runs through the park against the background of trees. The man's entire figure, face, arms and legs are visible. Anime style. The best quality.", ] examples_negative = [ "Blurred details, low resolution, no face visible, poor image of a man's face, poor quality, artifacts, black and white image.", ] css = """ #col-container { margin: 0 auto; max-width: 640px; } """ available_models = [ "stable-diffusion-v1-5/stable-diffusion-v1-5", "CompVis/stable-diffusion-v1-4", ] with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown(" # Text-to-Image Gradio Template from V. Gorsky") with gr.Row(): model = gr.Dropdown( label="Model Selection", choices=available_models, value="stable-diffusion-v1-5/stable-diffusion-v1-5", interactive=True ) prompt = gr.Textbox( label="Prompt", max_lines=1, placeholder="Enter your prompt", ) negative_prompt = gr.Textbox( label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", ) with gr.Row(): lora_scale = gr.Slider( label="LoRA scale", minimum=0.0, maximum=1.0, step=0.05, value=0.5, ) with gr.Row(): guidance_scale = gr.Slider( label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=7.5, ) with gr.Row(): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=4, ) with gr.Row(): num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=100, step=1, value=30, ) with gr.Accordion("Advanced Settings", open=False): with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512, ) with gr.Row(): height = gr.Slider( label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512, ) # ControlNet --------------------------------------------------------------------------------- with gr.Blocks(): with gr.Row(): use_control_net = gr.Checkbox( label="Use ControlNet", value=False, ) with gr.Column(visible=False) as control_net_options: strength_cn = gr.Slider( label="Strength", minimum=0.0, maximum=1.0, value=0.5, step=0.05, ) control_strength = gr.Slider( label="Control Strength", minimum=0.0, maximum=1.0, value=0.5, step=0.05, ) control_mode = gr.Dropdown( label="Control Mode", choices=[ "pose_estimation", ], value="pose_estimation", ) cn_source_image = gr.Image(label="Upload Source Image") control_image = gr.Image(label="Upload Control Net Image") use_control_net.change( fn=lambda x: gr.Row.update(visible=x), inputs=use_control_net, outputs=control_net_options ) # IP_adapter --------------------------------------------------------------------------------- with gr.Blocks(): with gr.Row(): use_ip_adapter = gr.Checkbox( label="Use IP_adapter", value=False, ) with gr.Column(visible=False) as ip_adapter_options: strength_ip = gr.Slider( label="Strength", minimum=0.0, maximum=1.0, value=0.5, step=0.05, ) ip_adapter_strength = gr.Slider( label="IP_adapter Strength", minimum=0.0, maximum=1.0, value=0.5, step=0.05, ) ip_source_image = gr.Image(label="Upload Source Image") ip_adapter_image = gr.Image(label="Upload IP_adapter Image") use_ip_adapter.change( fn=lambda x: gr.Row.update(visible=x), inputs=use_ip_adapter, outputs=ip_adapter_options ) # -------------------------------------------------------------------------------------- gr.Examples(examples=examples, inputs=[prompt]) gr.Examples(examples=examples_negative, inputs=[negative_prompt]) run_button = gr.Button("Run", scale=1, variant="primary") result = gr.Image(label="Result", show_label=False) gr.on( triggers=[run_button.click, prompt.submit], fn=infer, inputs=[ prompt, negative_prompt, width, height, num_inference_steps, model, seed, guidance_scale, lora_scale, strength_cn, # Коэфф. зашумления ControlNet use_control_net, # Чекбокс для ControlNet control_strength, # Контроль силы ControlNet cn_source_image, # Исходное изображение ControlNet control_image, # Контрольное изображение ControlNet strength_ip, # Коэфф. зашумления IP_adapter use_ip_adapter, # Параметр для включения IP_adapter ip_adapter_strength,# Сила влияния IP_adapter ip_source_image, # Исходное изображение IP_adapter ip_adapter_image, # Контрольное изображение IP_adapter ], outputs=[result], ) if __name__ == "__main__": demo.launch()