app.py

import gradio as gr
import numpy as np
import torch
from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline, StableDiffusionControlNetImg2ImgPipeline
from peft import PeftModel, LoraConfig
import os
from PIL import Image

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_default = "stable-diffusion-v1-5/stable-diffusion-v1-5"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# Инициализация ControlNet
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-openpose", torch_dtype=torch_dtype)

def get_lora_sd_pipeline(
    lora_dir='./lora_man_animestyle',
    base_model_name_or_path=None, 
    dtype=torch.float16, 
    adapter_name="default"
):
    unet_sub_dir = os.path.join(lora_dir, "unet")
    text_encoder_sub_dir = os.path.join(lora_dir, "text_encoder")
    
    if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:
        config = LoraConfig.from_pretrained(text_encoder_sub_dir)
        base_model_name_or_path = config.base_model_name_or_path
    
    if base_model_name_or_path is None:
        raise ValueError("Укажите название базовой модели или путь к ней")
    
    pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype)
    before_params = pipe.unet.parameters()
    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
    pipe.unet.set_adapter(adapter_name)
    after_params = pipe.unet.parameters()
    
    if os.path.exists(text_encoder_sub_dir):
        pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name)
    
    if dtype in (torch.float16, torch.bfloat16):
        pipe.unet.half()
        pipe.text_encoder.half()
    
    return pipe

def long_prompt_encoder(prompt, tokenizer, text_encoder, max_length=77):
    tokens = tokenizer(prompt, truncation=False, return_tensors="pt")["input_ids"]
    part_s = [tokens[:, i:i + max_length] for i in range(0, tokens.shape[1], max_length)]
    with torch.no_grad():
        embeds = [text_encoder(part.to(text_encoder.device))[0] for part in part_s]
    return torch.cat(embeds, dim=1)

def align_embeddings(prompt_embeds, negative_prompt_embeds):
    max_length = max(prompt_embeds.shape[1], negative_prompt_embeds.shape[1])
    return torch.nn.functional.pad(prompt_embeds, (0, 0, 0, max_length - prompt_embeds.shape[1])), \
           torch.nn.functional.pad(negative_prompt_embeds, (0, 0, 0, max_length - negative_prompt_embeds.shape[1]))

pipe_default = get_lora_sd_pipeline(lora_dir='./lora_man_animestyle', base_model_name_or_path=model_default, dtype=torch_dtype).to(device)
pipe_controlnet = StableDiffusionControlNetPipeline.from_pretrained(
    model_default,
    controlnet=controlnet,
    torch_dtype=torch_dtype
).to(device)

def preprocess_image(image, target_width, target_height):
    """
    Преобразует изображение в формат, подходящий для модели.
    """
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.resize((target_width, target_height), Image.LANCZOS)
    image = np.array(image).astype(np.float32) / 255.0  # Нормализация [0, 1]
    image = image[None].transpose(0, 3, 1, 2)  # Преобразуем в (batch, channels, height, width)
    image = torch.from_numpy(image).to(device)
    return image

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,   # Сила влияния ControlNet
    source_image=None,      # Исходное изображение
    control_image=None,     # Контрольное изображение
    progress=gr.Progress(track_tqdm=True)
):
    generator = torch.Generator(device).manual_seed(seed)
    
    if use_control_net and control_image is not None and source_image is not None:
        # Преобразуем изображения
        source_image = preprocess_image(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)
        
        # Используем ControlNet с LoRA
        pipe = pipe_controlnet
        prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder)
        negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder)
        prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
        image = pipe_controlnet(
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            image=source_image,     
            control_image=control_image,  
            width=width,
            height=height,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            controlnet_conditioning_scale=control_strength,
            generator=generator
        ).images[0]
    else:
        # Стандартная генерация без ControlNet
        if model != model_default:
            pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device)
            prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder)
            negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder)
            prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
        else:
            pipe = pipe_default
            prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder)
            negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder)
            prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
            pipe.fuse_lora(lora_scale=lora_scale)
        
        params = {
            'prompt_embeds': prompt_embeds,
            'negative_prompt_embeds': negative_prompt_embeds,
            'guidance_scale': guidance_scale,
            'num_inference_steps': num_inference_steps,
            'width': width,
            'height': height,
            'generator': generator,
        }
        
        image = pipe(**params).images[0]
    
    return image

examples = [
    "A young man in anime style. The image is characterized by high definition and resolution. Handsome, thoughtful man, attentive eyes. The man is depicted in the foreground, close-up or in the middle. High-quality images of the face, eyes, nose, lips, hands and clothes. The background and background are blurred and indistinct. The play of light and shadow is visible on the face and clothes.",
]    

examples_negative = [
    "blurred details, low resolution, poor image of a man's face, poor quality, artifacts, black and white image",
]

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

available_models = [
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    "CompVis/stable-diffusion-v1-4",
]

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # Text-to-Image Gradio Template from V. Gorsky")

        with gr.Row():
            model = gr.Dropdown(
                label="Model Selection",
                choices=available_models,
                value="stable-diffusion-v1-5/stable-diffusion-v1-5",
                interactive=True
            )
        
        prompt = gr.Textbox(
            label="Prompt",
            max_lines=1,
            placeholder="Enter your prompt",
        )

        negative_prompt = gr.Textbox(
            label="Negative prompt",
            max_lines=1,
            placeholder="Enter a negative prompt",
        )

        with gr.Row():
            lora_scale = gr.Slider(
                label="LoRA scale",
                minimum=0.0,
                maximum=1.0,
                step=0.05,
                value=0.5,
            ) 

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

        with gr.Row():
            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=1,
                maximum=100,
                step=1,
                value=30,
            )
               
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,
                )
            
            with gr.Row():
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,
                )

        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:    
                control_strength = gr.Slider(
                    label="Control Strength",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.5,
                    step=0.05,
                )
        
                control_mode = gr.Dropdown(
                    label="Control Mode",
                    choices=[
                        "pose_estimation",
                    ],
                    value="pose_estimation",
                )
        
                source_image = gr.Image(label="Upload Source Image")
                
                control_image = gr.Image(label="Upload Control Image")

            use_control_net.change(
                fn=lambda x: gr.Row.update(visible=x),
                inputs=use_control_net,
                outputs=control_net_options,
            )

        gr.Examples(examples=examples, inputs=[prompt])
        gr.Examples(examples=examples_negative, inputs=[negative_prompt])
        
        run_button = gr.Button("Run", scale=1, variant="primary")
        result = gr.Image(label="Result", show_label=False)   
    
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            width,
            height,
            num_inference_steps,
            model,
            seed,
            guidance_scale,
            lora_scale,
            use_control_net,   # Добавляем чекбокс для ControlNet
            control_strength,  # Добавляем контроль силы
            source_image,      # Добавляем исходное изображение 
            control_image,     # Добавляем контрольное изображение
        ],
        outputs=[result],
    )

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