Update app.py

app.py

@@ -1,74 +1,556 @@
 import gradio as gr
 import numpy as np
-import random
-from diffusers import DiffusionPipeline
 import torch
-
-###########################################
+from diffusers import (
+    StableDiffusionPipeline,
+    ControlNetModel,
+    StableDiffusionControlNetPipeline,
+    StableDiffusionControlNetImg2ImgPipeline,
+    AutoPipelineForImage2Image,
+    DDIMScheduler,
+    UniPCMultistepScheduler)
+from transformers import pipeline
+from diffusers.utils import load_image, make_image_grid
+from peft import PeftModel, LoraConfig
+import os
+from PIL import Image
 from rembg import remove 
-from PIL import Image 
-###############################################
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+IP_ADAPTER = 'h94/IP-Adapter'
+WEIGHT_NAME = "ip-adapter_sd15.bin"
+WEIGHT_NAME_plus = "ip-adapter-plus_sd15.bin"
+WEIGHT_NAME_face = "ip-adapter-full-face_sd15.bin"
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+model_default = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+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")
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+    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, resize_to_224=False):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    # Если resize_to_224=True, изменяем размер до 224x224
+    if resize_to_224:
+        image = image.resize((224, 224), Image.LANCZOS)
+    else:
+        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
+
+def get_depth_map(image, depth_estimator):
+    # Преобразуем изображение в PIL, если это необходимо
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    elif isinstance(image, torch.Tensor):
+        image = Image.fromarray(image.cpu().numpy())
+    # Получаем карту глубины
+    depth_map = depth_estimator(image)["depth"]
+    depth_map = np.array(depth_map)
+    depth_map = depth_map[:, :, None]  # Добавляем третье измерение
+    depth_map = np.concatenate([depth_map, depth_map, depth_map], axis=2)  # Преобразуем в 3 канала
+    depth_map = torch.from_numpy(depth_map).float() / 255.0  # Нормализация [0, 1]
+    depth_map = depth_map.permute(2, 0, 1)  # Меняем порядок осей (C, H, W)
+    return depth_map
+  
+pipe_default = get_lora_sd_pipeline(lora_dir='lora_man_animestyle', base_model_name_or_path=model_default, dtype=torch_dtype).to(device)
+
+# ----------------------------------------------------------------------------------------------------------------------------------------------------
 def infer(
-    model,
-    prompt,
-    negative_prompt,
-    seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    remove_bg, #################################################################
-    progress=gr.Progress(track_tqdm=True),
+    prompt, 
+    negative_prompt, 
+    model=model_default, 
+    width=512, 
+    height=512, 
+    num_inference_steps=50, 
+    seed=4, 
+    guidance_scale=7.5, 
+    lora_scale=0.7,
+    use_control_net=False,  # Параметр для включения ControlNet
+    control_mode=None,      # Параметр для выбора режима ControlNet
+    strength_cn=0.5,        # Коэфф. зашумления ControlNet
+    control_strength=0.5,   # Сила влияния ControlNet
+    cn_source_image=None,   # Исходное изображение ControlNet
+    control_image=None,     # Контрольное изображение ControlNet
+    use_ip_adapter=False,   # Параметр для включения IP_adapter
+    ip_adapter_mode=None,   # Параметр для выбора режима IP_adapter
+    strength_ip=0.5,        # Коэфф. зашумления IP_adapter
+    ip_adapter_strength=0.5,# Сила влияния IP_adapter
+    controlnet_conditioning_scale=0.5, # Сила влияния ControlNet
+    ip_source_image=None,   # Исходное изображение IP_adapter
+    ip_adapter_image=None,  # Контрольное изображение IP_adapter
+    progress=gr.Progress(track_tqdm=True)
 ):
-    global model_repo_id
-    if model != model_repo_id:
-        print(model, model_repo_id)
-        pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype)
-        pipe = pipe.to(device)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    #############################################################
-    # Если выбрано удаление фона
-    if remove_bg:
-        image = remove(image)
-    ##############################################################
+
+    # Генерация изображений с Ip_Adapter ------------------------------------------------------------------------------------------------------------------
+    if use_ip_adapter and ip_source_image is not None and ip_adapter_image is not None:
+        
+        if ip_adapter_mode == "pose_estimation":    
+
+            print('ip_adapter_mode = ', ip_adapter_mode)
+            
+            # Инициализация ControlNet
+            controlnet_model_path = "lllyasviel/sd-controlnet-openpose"
+            controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype)
+
+            generator = torch.Generator(device).manual_seed(seed)
+
+            pipe_ip_adapter = StableDiffusionControlNetPipeline.from_pretrained(
+                model_default,
+                controlnet=controlnet,
+                torch_dtype=torch_dtype
+            ).to(device)
+
+            # Загрузка IP-Adapter
+            pipe_ip_adapter.load_ip_adapter(IP_ADAPTER, subfolder="models", weight_name=WEIGHT_NAME_plus)
+            pipe_ip_adapter.set_ip_adapter_scale(ip_adapter_strength)
+
+            # Преобразование изображений для IP-Adapter (размер 224x224)
+            ip_source_image = preprocess_image(ip_source_image, width, height, resize_to_224=True)
+            ip_adapter_image = preprocess_image(ip_adapter_image, width, height, resize_to_224=True)
+
+            # Создаём пайплайн IP_adapter с LoRA, если он ещё не создан
+            if not hasattr(pipe_ip_adapter, 'lora_loaded') or not pipe_ip_adapter.lora_loaded:
+                # Загружаем LoRA для UNet
+                pipe_ip_adapter.unet = PeftModel.from_pretrained(
+                    pipe_ip_adapter.unet,
+                    'lora_man_animestyle/unet',
+                    adapter_name="default"
+                )
+                pipe_ip_adapter.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_ip_adapter.text_encoder = PeftModel.from_pretrained(
+                        pipe_ip_adapter.text_encoder,
+                        text_encoder_lora_path,
+                        adapter_name="default"
+                    )
+                    pipe_ip_adapter.text_encoder.set_adapter("default")
+
+                # Объединяем LoRA с основной моделью
+                pipe_ip_adapter.fuse_lora(lora_scale=lora_scale)
+                pipe_ip_adapter.lora_loaded = True  # Помечаем, что LoRA загружена
+                
+            # Убедимся, что параметры имеют тип float        
+            ip_adapter_strength = float(ip_adapter_strength)
+            controlnet_conditioning_scale = float(controlnet_conditioning_scale)
+
+            # Используем IP-Adapter с LoRA
+            prompt_embeds = long_prompt_encoder(prompt, pipe_ip_adapter.tokenizer, pipe_ip_adapter.text_encoder)
+            negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe_ip_adapter.tokenizer, pipe_ip_adapter.text_encoder)
+            prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
+            image = pipe_ip_adapter(
+                prompt_embeds=prompt_embeds,
+                negative_prompt_embeds=negative_prompt_embeds,
+                image=ip_adapter_image, 
+                ip_adapter_image=ip_source_image, 
+                strength=strength_ip, 
+                width=width,
+                height=height,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                controlnet_conditioning_scale=controlnet_conditioning_scale,
+                generator=generator,
+            ).images[0]
+        else:
+
+            if ip_adapter_mode == "edge_detection":    
+
+                print('ip_adapter_mode = ', ip_adapter_mode)
+
+                # Инициализация ControlNet
+                controlnet_model_path = "lllyasviel/control_v11f1p_sd15_depth"
+                controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype)
+
+                generator = torch.Generator(device).manual_seed(seed)
+
+                pipe_ip_adapter = StableDiffusionControlNetPipeline.from_pretrained(
+                    model_default,
+                    controlnet=controlnet,
+                    torch_dtype=torch_dtype
+                ).to(device)
+
+                # Загрузка IP-Adapter
+                #pipe_ip_adapter.load_ip_adapter(IP_ADAPTER, subfolder="models", weight_name=WEIGHT_NAME_face)
+                pipe_ip_adapter.load_ip_adapter(IP_ADAPTER, subfolder="models", weight_name=WEIGHT_NAME_plus)
+                pipe_ip_adapter.set_ip_adapter_scale(ip_adapter_strength)
+
+                # Преобразование изображений для IP-Adapter (размер 224x224)
+                ip_source_image = preprocess_image(ip_source_image, width, height, resize_to_224=True)
+                ip_adapter_image = preprocess_image(ip_adapter_image, width, height, resize_to_224=True)
+
+                # Создаём пайплайн IP_adapter с LoRA, если он ещё не создан
+                if not hasattr(pipe_ip_adapter, 'lora_loaded') or not pipe_ip_adapter.lora_loaded:
+                    # Загружаем LoRA для UNet
+                    pipe_ip_adapter.unet = PeftModel.from_pretrained(
+                        pipe_ip_adapter.unet,
+                        'lora_man_animestyle/unet',
+                        adapter_name="default"
+                    )
+                    pipe_ip_adapter.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_ip_adapter.text_encoder = PeftModel.from_pretrained(
+                            pipe_ip_adapter.text_encoder,
+                            text_encoder_lora_path,
+                            adapter_name="default"
+                        )
+                        pipe_ip_adapter.text_encoder.set_adapter("default")
+
+                    # Объединяем LoRA с основной моделью            
+                    pipe_ip_adapter.fuse_lora(lora_scale=lora_scale)
+                    pipe_ip_adapter.lora_loaded = True  # Помечаем, что LoRA загружена
+
+                # Убедимся, что параметры имеют тип float        
+                ip_adapter_strength = float(ip_adapter_strength)
+                controlnet_conditioning_scale = float(controlnet_conditioning_scale)
+
+                # Используем IP-Adapter с LoRA
+                prompt_embeds = long_prompt_encoder(prompt, pipe_ip_adapter.tokenizer, pipe_ip_adapter.text_encoder)
+                negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe_ip_adapter.tokenizer, pipe_ip_adapter.text_encoder)
+                prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
+
+                # scale = { # по умолчанию в остальных блоках везде 0.
+                #    "down": {
+                #         "block_0": [0.0, 1.0],
+                #         "block_1": [0.0, 1.0],
+                #     },
+                #     "up": {
+                #         "block_0": [0.0, 1.0, 0.0],
+                #         "block_1": [0.0, 1.0, 0.0],
+                #     },
+                # }
+                # scale = {
+                # "down": {"block_2": [0.0, 1.0]},
+                # "up": {"block_0": [0.0, 1.0, 0.0]},
+                # }
+                # pipe_ip_adapter.set_ip_adapter_scale(scale)
+
+                image = pipe_ip_adapter(
+                    prompt_embeds=prompt_embeds,
+                    negative_prompt_embeds=negative_prompt_embeds,
+                    image=ip_adapter_image, 
+                    ip_adapter_image=ip_source_image, 
+                    strength=strength_ip, 
+                    width=width,
+                    height=height,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=guidance_scale,
+                    controlnet_conditioning_scale=controlnet_conditioning_scale,
+                    generator=generator,
+                ).images[0]    
+            else:
+
+                if ip_adapter_mode == "depth_map":    
+
+                    print('ip_adapter_mode = ', ip_adapter_mode)
+
+                    # Убедимся, что параметры имеют тип float        
+                    controlnet_conditioning_scale = float(controlnet_conditioning_scale)
+                    
+                    # Инициализация ControlNet
+                    controlnet_model_path = "lllyasviel/control_v11f1p_sd15_depth"
+                    controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype)
+
+                    generator = torch.Generator(device).manual_seed(seed)
+
+                    # Преобразование изображений для IP-Adapter (размер 224x224)
+                    ip_source_image = preprocess_image(ip_source_image, width, height, resize_to_224=True)
+                    ip_adapter_image = preprocess_image(ip_adapter_image, width, height, resize_to_224=True)
+
+                    pipe_ip_adapter = StableDiffusionControlNetPipeline.from_pretrained(
+                            model_default,
+                            controlnet=controlnet,
+                            torch_dtype=torch_dtype
+                        ).to(device)
+                    pipe_ip_adapter.load_ip_adapter(IP_ADAPTER, subfolder="models", weight_name=WEIGHT_NAME)
+
+                    pipe_ip_adapter.set_ip_adapter_scale(ip_adapter_strength)
+                    image = pipe_ip_adapter(
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        image=ip_source_image,
+                        width=width,
+                        height=height,
+                        ip_adapter_image=ip_adapter_image,
+                        num_inference_steps=num_inference_steps,
+                        strength=strength_ip, 
+                        guidance_scale=guidance_scale,
+                        controlnet_conditioning_scale=controlnet_conditioning_scale,
+                        generator=generator,
+                    ).images[0]
+                else:    
+
+                    if ip_adapter_mode == "face_model":     
+
+                        print('ip_adapter_mode = ', ip_adapter_mode)
+
+                        # Преобразование изображений для IP-Adapter (размер 224x224)
+                        ip_source_image = preprocess_image(ip_source_image, width, height, resize_to_224=True)
+                        ip_adapter_image = preprocess_image(ip_adapter_image, width, height, resize_to_224=True)
+                        
+                        pipe_ip_adapter = StableDiffusionPipeline.from_pretrained(
+                            model_default,
+                            torch_dtype=torch_dtype,
+                        ).to(device)
+
+                        pipe_ip_adapter.scheduler = DDIMScheduler.from_config(pipe_ip_adapter.scheduler.config)
+                        pipe_ip_adapter.load_ip_adapter(IP_ADAPTER, subfolder="models", weight_name=WEIGHT_NAME_face)
+
+                        generator = torch.Generator(device).manual_seed(seed)
+
+                        pipe_ip_adapter.set_ip_adapter_scale(ip_adapter_strength)
+                        image = pipe_ip_adapter(
+                            prompt=prompt,
+                            negative_prompt=negative_prompt,
+                            ip_adapter_image=ip_adapter_image,
+                            width=width,
+                            height=height,
+                            guidance_scale=guidance_scale,
+                            num_inference_steps=num_inference_steps,
+                            generator=generator,
+                        ).images[0]
+    else:
+        # Генерация изображений с ControlNet ----------------------------------------------------------------------------------------------------------------
+
+        if use_control_net and control_image is not None and cn_source_image is not None:
+
+            if control_mode == "pose_estimation":     
+
+                print('control_mode = ', control_mode)
+
+                # Инициализация ControlNet
+                controlnet_model_path = "lllyasviel/sd-controlnet-openpose"
+                controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype)
+
+                generator = torch.Generator(device).manual_seed(seed)
+
+                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
+                prompt_embeds = long_prompt_encoder(prompt, pipe_controlnet.tokenizer, pipe_controlnet.text_encoder)
+                negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe_controlnet.tokenizer, pipe_controlnet.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:
+
+                if control_mode == "edge_detection":     
+
+                    print('control_mode = ', control_mode)
+
+                    controlnet_model_path = "lllyasviel/sd-controlnet-canny"
+                    controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype, use_safetensors=True)
+
+                    generator = torch.Generator(device).manual_seed(seed)
+
+                    pipe_controlnet = StableDiffusionControlNetPipeline.from_pretrained(
+                        model_default, 
+                        controlnet=controlnet, 
+                        torch_dtype=torch_dtype, 
+                        use_safetensors=True
+                    ).to(device)
+
+                    pipe_controlnet.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)
+
+                    # Преобразуем изображения
+                    cn_source_image = preprocess_image(cn_source_image, width, height)
+                    control_image = preprocess_image(control_image, width, height)
+
+                    image = pipe_controlnet(
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        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:
+
+                    if control_mode == "depth_map":     
+
+                        print('control_mode = ', control_mode)
+
+                        depth_estimator = pipeline("depth-estimation")
+                        depth_map = get_depth_map(control_image, depth_estimator).unsqueeze(0).half().to(device)
+
+                        controlnet_model_path = "lllyasviel/control_v11f1p_sd15_depth"
+                        controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype, use_safetensors=True)
+
+                        generator = torch.Generator(device).manual_seed(seed)
+
+                        pipe_controlnet = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+                            model_default, 
+                            controlnet=controlnet,
+                            torch_dtype=torch_dtype,
+                            use_safetensors=True
+                        ).to(device)
+
+                        pipe_controlnet.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)
     
-    return image, seed
+                        image = pipe_controlnet(
+                            prompt=prompt,
+                            negative_prompt=negative_prompt,
+                            image=control_image,
+                            control_image=depth_map,
+                            width=width,
+                            height=height,
+                            num_inference_steps=num_inference_steps,
+                            guidance_scale=guidance_scale,
+                            generator=generator
+                        ).images[0]
+        else:
+            # Генерация изображений с LORA без ControlNet и IP_Adapter ---------------------------------------------------------------------------------------------
+
+            # Инициализация ControlNet
+            controlnet_model_path = "lllyasviel/sd-controlnet-openpose"
+            controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype)
+
+            generator = torch.Generator(device).manual_seed(seed)
 
+            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 = [
-    "Young man in anime style. The image is of high sharpness and resolution. A handsome, thoughtful man. The man is depicted in the foreground, close-up or middle plan. The background is blurry, not sharp. The play of light and shadow is visible on the face and clothes."
-]
+    "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.",
+    "The smiling man. His face and hands are visible. Anime style. The best quality.",
+    "The smiling girl. Anime style. Best quality, high quality.",
+    "lego batman and robin. Rich and vibrant colors.",
+    "A photo of Pushkin as a hockey player in uniform with a stick, playing hockey on the ice arena in the NHL and scoring a goal.",
+]    
 
 examples_negative = [
-    "blurred details, low resolution, poor image of a man's face, poor quality, artifacts, black and white image"
+    "Blurred details, low resolution, bad anatomy, no face visible, poor image of a man's face, poor quality, artifacts, black and white image.",
+    "Monochrome, lowres, bad anatomy, worst quality, low quality",
+    "lowres, bad anatomy, worst quality, low quality, black and white image.",
 ]
 
 css = """
@@ -83,54 +565,67 @@ available_models = [
     "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")
 
-        model = gr.Dropdown(
-            label="Model Selection",
-            choices=available_models,
-            value="stable-diffusion-v1-5/stable-diffusion-v1-5",
-            interactive=True
-        )
-        prompt = gr.Text(
+        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",
-            show_label=False,
             max_lines=1,
             placeholder="Enter your prompt",
-            container=False,
         )
 
-        negative_prompt = gr.Text(
+        negative_prompt = gr.Textbox(
             label="Negative prompt",
             max_lines=1,
             placeholder="Enter a negative prompt",
-            visible=True,
         )
+
+        with gr.Row():
+            lora_scale = gr.Slider(
+                label="LoRA scale",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.7,
+            ) 
+
+        with gr.Row():
+            guidance_scale = gr.Slider(
+                label="Guidance scale",
+                minimum=0.0,
+                maximum=10.0,
+                step=0.01,
+                value=7.5,
+            )  
         
-        seed = gr.Slider(
+        with gr.Row():
+            seed = gr.Slider(
                 label="Seed",
                 minimum=0,
                 maximum=MAX_SEED,
                 step=1,
-                value=0,
-        )
-        guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=7.5,  
-        )
-        num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=100,
-                    step=1,
-                    value=30,  
-        )
+                value=4,
+            )
 
+        with gr.Row():
+            num_inference_steps = gr.Slider(
+                label="Number of inference steps",
+                minimum=1,
+                maximum=100,
+                step=1,
+                value=50,
+            )
+               
         with gr.Accordion("Advanced Settings", open=False):
             with gr.Row():
                 width = gr.Slider(
@@ -138,49 +633,160 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,  
+                    value=512,
                 )
-
+            
+            with gr.Row():
                 height = gr.Slider(
                     label="Height",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,  
+                    value=512,
                 )
 
-        #########################################################
-        # Добавляем Checkbox для удаления фона
-        remove_bg = gr.Checkbox(
-            label="Remove Background",
-            value=False,
-            interactive=True
-        )
-        ##########################################################
+        # ControlNet -----------------------------------------------------------------------------------------------
+        with gr.Blocks():
+            with gr.Row():
+                use_control_net = gr.Checkbox(
+                    label="Use ControlNet",
+                    value=False,
+                )
 
-        gr.Examples(examples=examples, inputs=[prompt])
-        gr.Examples(examples=examples_negative, inputs=[negative_prompt])
+            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.01,
+                    interactive=True,
+                )
+                
+                control_strength = gr.Slider(
+                    label="Control Strength",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.5,
+                    step=0.01,
+                    interactive=True,
+                )
+
+                control_mode = gr.Dropdown(
+                    label="Control Mode",
+                    choices=[
+                        "pose_estimation",
+                        "edge_detection", 
+                        "depth_map",
+                    ],
+                    value="pose_estimation",
+                    interactive=True,
+                )
         
-        run_button = gr.Button("Run", scale=0, variant="primary")
-        result = gr.Image(label="Result", show_label=False)
+                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.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.01,
+                    interactive=True,
+                )
+                
+                ip_adapter_strength = gr.Slider(
+                    label="IP_Adapter Strength",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.5,
+                    step=0.01,
+                    interactive=True,
+                )
+
+                controlnet_conditioning_scale = gr.Slider(
+                    label="Controlnet conditioning scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.5,
+                    step=0.01,
+                    interactive=True,
+                )
+
+                ip_adapter_mode = gr.Dropdown(
+                    label="Ip_Adapter Mode",
+                    choices=[
+                        "pose_estimation",
+                        "edge_detection", 
+                        "depth_map",
+                        "face_model"
+                    ],
+                    value="pose_estimation",
+                    interactive=True,
+                )
+
+                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.update(visible=x), 
+                inputs=use_ip_adapter,
+                outputs=ip_adapter_options
+            )
+        # ---------------------------------------------------------------------------------------------------------
         
+        gr.Examples(examples=examples, inputs=[prompt], label="Examples for prompt:")
+        gr.Examples(examples=examples_negative, inputs=[negative_prompt], label="Examples for 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=[
-            model,
             prompt,
             negative_prompt,
-            seed,
+            model,
             width,
             height,
-            guidance_scale,
             num_inference_steps,
-            remove_bg,  # Добавляем remove_bg в inputs ###############################################
+            seed,
+            guidance_scale,
+            lora_scale,
+            use_control_net,    # Параметр для включения ControlNet
+            control_mode,       # Параметр для выбора режима ControlNet
+            strength_cn,        # Коэфф. зашумления ControlNet
+            control_strength,   # Сила влияния ControlNet
+            cn_source_image,    # Исходное изображение ControlNet
+            control_image,      # Контрольное изображение ControlNet
+            use_ip_adapter,     # Параметр для в��лючения IP_adapter
+            ip_adapter_mode,    # Параметр для выбора режима IP_adapter
+            strength_ip,        # Коэфф. зашумления IP_adapter
+            ip_adapter_strength,# Сила влияния IP_adapter
+            controlnet_conditioning_scale, # Сила влияния ControlNet
+            ip_source_image,    # Исходное изображение IP_adapter
+            ip_adapter_image,   # Контрольное изображение IP_adapter
         ],
-        outputs=[result, seed],
+        outputs=[result],
     )
 
 if __name__ == "__main__":
     demo.launch()
-