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")
    
    # Проверка существования директорий LoRA
    print(f"LoRA directory exists: {os.path.exists(lora_dir)}")
    print(f"UNet LoRA exists: {os.path.exists(unet_sub_dir)}")
    print(f"Text encoder LoRA exists: {os.path.exists(text_encoder_sub_dir)}")
    
    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)
    
    # Логирование параметров до применения LoRA
    before_params = list(pipe.unet.parameters())
    
    # Применение LoRA к UNet
    if os.path.exists(unet_sub_dir):
        pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
        pipe.unet.set_adapter(adapter_name)
    
    # Применение LoRA к текстовому энкодеру (если есть)
    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)
    
    # Логирование параметров после применения LoRA
    after_params = list(pipe.unet.parameters())
    print(f"Parameters changed: {before_params != after_params}")
    
    # Детальное сравнение параметров
    for i, (param1, param2) in enumerate(zip(before_params, after_params)):
        if not torch.equal(param1, param2):
            print(f"Parameter {i} changed.")
        else:
            print(f"Parameter {i} did not change.")
    
    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,
    progress=gr.Progress(track_tqdm=True)
):
    print(f"Received lora_scale: {lora_scale}")  # Лог для проверки значения lora_scale
    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)
        
        # Логирование параметров до и после применения LoRA
        before_params = list(pipe.unet.parameters())
        print(f"Applying LoRA with scale: {lora_scale}")
        pipe.fuse_lora(lora_scale=lora_scale)
        after_params = list(pipe.unet.parameters())
        print(f"Parameters changed: {before_params != after_params}")
    
    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,
    }
    
    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.",
]    

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

        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,
        ],
        outputs=[result],
    )

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