app.py

import gradio as gr
import numpy as np
import torch
from diffusers import (
    StableDiffusionPipeline,
    ControlNetModel,
    StableDiffusionControlNetPipeline,
    StableDiffusionControlNetImg2ImgPipeline,
    AutoPipelineForImage2Image,
    DDIMScheduler,
    UniPCMultistepScheduler,
    LCMScheduler, 
    AutoPipelineForText2Image,
    DPMSolverMultistepScheduler,
    AutoencoderKL,
)
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 

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"

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

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, 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(
    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
    remove_bg=None,          # Удаление фона с изображения
    use_LCM_adapter=False,   # Параметр для включения LCM_adapter
    LCM_adapter=None,        # Параметр для выбора типа LCM_adapter
    use_DDIMScheduler=False, # Параметр для включения DDIMScheduler
    use_Tiny_VAE=False,      # Параметр для включения Tiny_VAE
    Tiny_VAE=None,           # Параметр для выбора типа Tiny_VAE
    progress=gr.Progress(track_tqdm=True)
):

    # Генерация изображений с 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)

                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)
    
                        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:
            # Генерация изображений с LCM_Adapter ---------------------------------------------------------------------------------------------
            
            if use_LCM_adapter:

                print('use_LCM_adapter = ', use_LCM_adapter)
                
                if LCM_adapter == "lcm-lora-sdv1-5":
                    adapter_id = "latent-consistency/lcm-lora-sdv1-5"

                guidance_scale = 0
                generator = torch.Generator(device).manual_seed(seed)

                pipe_LCM = get_lora_sd_pipeline(lora_dir='lora_man_animestyle', base_model_name_or_path=model_default, dtype=torch_dtype).to(device)
                pipe_LCM.scheduler = LCMScheduler.from_config(pipe_LCM.scheduler.config)
                pipe_LCM.to(device)

                pipe_LCM.load_lora_weights(adapter_id)
                pipe_LCM.fuse_lora()

                prompt_embeds = long_prompt_encoder(prompt, pipe_LCM.tokenizer, pipe_LCM.text_encoder)
                negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe_LCM.tokenizer, pipe_LCM.text_encoder)
                prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)

                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,
                    'cross_attention_kwargs': {"scale": lora_scale},
                }

                image = pipe_LCM(**params).images[0]    
            else:        
                # Генерация изображений с DDIMScheduler ---------------------------------------------------------------------------------------------

                if use_DDIMScheduler:

                    print('use_DDIMScheduler = ', use_DDIMScheduler)

                    generator = torch.Generator(device).manual_seed(seed)

                    pipe_DDIMS = StableDiffusionPipeline.from_pretrained(model_default, torch_dtype=torch_dtype).to(device)
                    pipe_DDIMS.scheduler = DPMSolverMultistepScheduler.from_config(pipe_DDIMS.scheduler.config)

                    prompt_embeds = long_prompt_encoder(prompt, pipe_DDIMS.tokenizer, pipe_DDIMS.text_encoder)
                    negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe_DDIMS.tokenizer, pipe_DDIMS.text_encoder)
                    prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)

                    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_DDIMS(**params).images[0]
                else:    
                    # Генерация изображений с Tiny_VAE ---------------------------------------------------------------------------------------------

                    if use_Tiny_VAE:
                        
                        print('use_Tiny_VAE = ', use_Tiny_VAE)

                        if Tiny_VAE == "sd-vae-ft-mse":
                            VAE_id = "stabilityai/sd-vae-ft-mse"
                        
                        generator = torch.Generator(device).manual_seed(seed)

                        vae = AutoencoderKL.from_pretrained(VAE_id, torch_dtype=torch_dtype)

                        pipe_Tiny_VAE = StableDiffusionPipeline.from_pretrained(model_default, vae=vae, torch_dtype=torch_dtype).to(device)

                        prompt_embeds = long_prompt_encoder(prompt, pipe_Tiny_VAE.tokenizer, pipe_Tiny_VAE.text_encoder)
                        negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe_Tiny_VAE.tokenizer, pipe_Tiny_VAE.text_encoder)
                        prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)

                        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_Tiny_VAE(**params).images[0]
                    else:    
                        # Генерация изображений с LORA без ControlNet и IP_Adapter ---------------------------------------------------------------------------------------------

                        print('Генерация изображений с 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_default = get_lora_sd_pipeline(lora_dir='lora_man_animestyle', base_model_name_or_path=model_default, dtype=torch_dtype).to(device)
                            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,
                            'cross_attention_kwargs': {"scale": lora_scale},
                        }

                        image = pipe(**params).images[0]

    # Если выбрано удаление фона
    if remove_bg:
        image = remove(image)            
    
    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.",
    "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, 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 = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

available_models = [
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    "nota-ai/bk-sdm-small",
    "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.01,
                value=0.7,
                interactive=True
            ) 

        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance scale",
                minimum=0.0,
                maximum=10.0,
                step=0.01,
                value=7.5,
                interactive=True
            )  
        
        with gr.Row():
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=4,
                interactive=True
            )

        with gr.Row():
            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=1,
                maximum=100,
                step=1,
                value=50,
                interactive=True
            )
               
        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.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,
                )
        
                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
            )

        # LCM_Adapter ------------------------------------------------------------------------------------------------
        with gr.Blocks():
            with gr.Row():
                use_LCM_adapter = gr.Checkbox(
                    label="Use LCM_Adapter",
                    value=False,
                    interactive=True
                )

                LCM_adapter = gr.Dropdown(
                    label="LCM_Adapter Selection",
                    choices=[
                        "lcm-lora-sdv1-5",
                    ],
                    value="lcm-lora-sdv1-5",
                    visible=False,
                    interactive=True,
                )

            use_LCM_adapter.change(
                fn=lambda x: gr.update(visible=x), 
                inputs=use_LCM_adapter,
                outputs=LCM_adapter
            )
        
        # DDIMScheduler ------------------------------------------------------------------------------------------------
        # Checkbox для DDIMScheduler
        with gr.Blocks():
            use_DDIMScheduler = gr.Checkbox(
                label="Use DDIMScheduler",
                value=False,
                interactive=True
            )

        # Tiny_VAE -----------------------------------------------------------------------------------------------------
        # Checkbox для Tiny_VAE
        with gr.Blocks():
            with gr.Row():
                use_Tiny_VAE = gr.Checkbox(
                    label="Use Tiny_VAE",
                    value=False,
                    interactive=True
                )

                Tiny_VAE = gr.Dropdown(
                    label="Tiny_VAE Selection",
                    choices=[
                        "sd-vae-ft-mse",
                    ],
                    value="sd-vae-ft-mse",
                    visible=False,
                    interactive=True,
                )

            use_Tiny_VAE.change(
                fn=lambda x: gr.update(visible=x), 
                inputs=use_Tiny_VAE,
                outputs=Tiny_VAE
            )
            
        # Удаление фона------------------------------------------------------------------------------------------------
        # Checkbox для удаления фона
        with gr.Blocks():
            remove_bg = gr.Checkbox(
                label="Remove Background",
                value=False,
                interactive=True
            )
        # -------------------------------------------------------------------------------------------------------------

        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=[
            prompt,
            negative_prompt,
            model,
            width,
            height,
            num_inference_steps,
            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
            remove_bg,          # Удаление фона с изображения
            use_LCM_adapter,    # Параметр для включения LCM_adapter
            LCM_adapter,        # Параметр для выбора типа LCM_adapter
            use_DDIMScheduler,  # Параметр для включения DDIMScheduler
            use_Tiny_VAE,       # Параметр для включения Tiny_VAE
            Tiny_VAE,           # Параметр для выбора типа Tiny_VAE
        ],
        outputs=[result],
    )

if __name__ == "__main__":
    demo.launch()