app.py

import gradio as gr
import numpy as np
import torch
from diffusers import StableDiffusionPipeline
from peft import PeftModel, LoraConfig
import os

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

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]))

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, 
    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,
    # use_control_net=False,  # Добавляем параметр для управления включением ControlNet
    # control_strength=0.5,  
    # use_ip_adapter=False,  # Добавляем параметр для управления включением Ip_Adapter
    # ip_adapter_scale=0.5,
    progress=gr.Progress(track_tqdm=True)
    ):
    
    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) # Коэфф. добавления lora
    
    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,
    }
    
    # if use_control_net:  # Если ControlNet включен
    #     params['control_net'] = True  # Включаем использование ControlNet
    #     params['control_strength'] = control_strength  # Устанавливаем вес ControlNet

    # if use_ip_adapter:  # Если Ip_Adapter включен
    #     params['ip_adapter'] = True  # Включаем использование Ip_Adapter
    #     params['ip_adapter_scale'] = ip_adapter_scale  # Устанавливаем вес Ip_Adapter    
    
    return pipe(**params).images[0]

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.",
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k.",
    "An astronaut riding a green horse.",
]    

examples_negative = [
    "blurred details, low resolution, poor image of a man's face, poor quality, artifacts, black and white image",
    "blurry details, low resolution, poorly defined edges",
    "bad face, bad quality, artifacts, low-res, black and white",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

available_models = [
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    "SG161222/Realistic_Vision_V3.0_VAE",
    "CompVis/stable-diffusion-v1-4",
    "stabilityai/sdxl-turbo",
    "runwayml/stable-diffusion-v1-5",
    "sd-legacy/stable-diffusion-v1-5",
    "prompthero/openjourney",
    "stabilityai/stable-diffusion-3-medium-diffusers",
    "stabilityai/stable-diffusion-3.5-large",
    "stabilityai/stable-diffusion-3.5-large-turbo",
]

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 ---------------------------------------------------------------------
        def process_input_ControlNet(image, use_control_net, control_strength, control_mode):
            if use_control_net:
                # Логика для обработки с использованием ControlNet
                result = f"ControlNet активен! Режим: {control_mode}, Интенсивность: {control_strength}"
            else:
                # Логика для обработки без ControlNet
                result = "ControlNet отключен."
            return result

        with gr.Blocks():
            with gr.Row():
                # Чекбокс для включения/отключения ControlNet
                use_control_net = gr.Checkbox(
                    label="Use ControlNet",
                    value=False,
                )

            # Дополнительные опции для ControlNet
            with gr.Column(visible=False) as control_net_options:    
                # Слайдер для настройки интенсивности
                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=[
                        "edge_detection",
                        "canny_edge_detection",
                        "pose_estimation",
                        "depth_map",
                        "segmentation_map",
                        "scribble_sketch",
                        "normal_map",
                        "hed_edge_detection",
                        "openpose",
                        "mlsd_line_detection",
                        "scribble_diffusion",
                        "semantic_segmentation",
                        "style_transfer",
                        "colorization",
                        "custom_map"
                    ],
                    value="pose_estimation",
                )
        
                # Окно для загрузки изображений
                control_image = gr.Image(label="Upload Control Image")

                # Кнопка для запуска работы ControlNet
                run_button = gr.Button("Run")

                # Текстовое поле для вывода результата
                output = gr.Textbox(label="Output")

            # Логика для отображения/скрытия дополнительных опций
            use_control_net.change(
                fn=lambda x: gr.Row.update(visible=x),
                inputs=use_control_net,
                outputs=control_net_options,
            )

            # Привязка кнопки Run к функции работы с ControlNet
            run_button.click(
                fn=process_input_ControlNet,
                inputs=[control_image, use_control_net, control_strength, control_mode],
                outputs=output,
            )

        # Функция для работы с IP-adapter ----------------------------------------------------------------------------
        def process_input_IP_adapter(image, use_ip_adapter, ip_adapter_scale, ip_adapter_image):
            if use_ip_adapter:
                # Логика для обработки с использованием IP-adapter
                result = f"IP-adapter активен! Масштаб: {ip_adapter_scale}"
            else:
                # Логика для обработки без IP-adapter
                result = "IP-adapter отключен."
            return result

        # Создание интерфейса
        with gr.Blocks():
            with gr.Row():
                # Чекбокс для включения/отключения IP-adapter
                use_ip_adapter = gr.Checkbox(
                    label="Use IP-adapter",
                    value=False,
                )

            # Дополнительные опции для IP-adapter
            with gr.Column(visible=False) as ip_adapter_options:
                # Слайдер для настройки масштаба
                ip_adapter_scale = gr.Slider(
                     label="IP-adapter Scale",
                     minimum=0.0,
                     maximum=1.0,
                     value=0.5,
                     step=0.05,
                )
        
                # Окно для загрузки изображений
                ip_adapter_image = gr.Image(label="Upload IP-adapter Image")

                # Кнопка для запуска обработки
                run_button = gr.Button("Run")

                # Текстовое поле для вывода результата
                output = gr.Textbox(label="Output")

            # Логика для отображения/скрытия дополнительных опций
            use_ip_adapter.change(
               fn=lambda x: gr.Row.update(visible=x),
                inputs=use_ip_adapter,
                outputs=ip_adapter_options,
            )

            # Привязка кнопки Run к функции работы с IP-adapter
            run_button.click(
                fn=process_input_IP_adapter,
                inputs=[ip_adapter_image, use_ip_adapter, ip_adapter_scale, ip_adapter_image],
                outputs=output,
            )

        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,
            # use_control_net,  
            # control_strength,  
            # use_ip_adapter,  
            # ip_adapter_scale,
        ],
        outputs=[result],
    )

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