Update app.py

Добавлена возможность использовать LoRA и модель по умолчанию заменена на CompVis/stable-diffusion-v1-5

app.py

@@ -6,10 +6,63 @@ import random
 from diffusers import DiffusionPipeline
 import torch
 
+from peft import PeftModel, LoraConfig
+import os
+
+def get_lora_sd_pipeline(
+    ckpt_dir='./lora_logos', 
+    base_model_name_or_path=None, 
+    dtype=torch.float16, 
+    adapter_name="default"
+    ):
+
+    unet_sub_dir = os.path.join(ckpt_dir, "unet")
+    text_encoder_sub_dir = os.path.join(ckpt_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("Please specify the base model name or path")
+    
+    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()
+    print("Parameters changed:", any(torch.any(b != a) for b, a in zip(before_params, after_params)))
+    
+    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 process_prompt(prompt, tokenizer, text_encoder, max_length=77):
+    tokens = tokenizer(prompt, truncation=False, return_tensors="pt")["input_ids"]
+    chunks = [tokens[:, i:i + max_length] for i in range(0, tokens.shape[1], max_length)]
+    
+    with torch.no_grad():
+        embeds = [text_encoder(chunk.to(text_encoder.device))[0] for chunk in chunks]
+    
+    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]))
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 #model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-model_repo_id = "CompVis/stable-diffusion-v1-4"
-model_dropdown = ['stabilityai/sdxl-turbo', 'CompVis/stable-diffusion-v1-4' ]
+model_id_default = "CompVis/stable-diffusion-v1-5"
+model_dropdown = ['stabilityai/sdxl-turbo', 'CompVis/stable-diffusion-v1-4', 'CompVis/stable-diffusion-v1-5'  ]
+
+model_lora_default = "lora_lady_and_cats_logos"
+model_lora_dropdown = ['lora_lady_and_cats_logos', 'lora_pussinboots_logos'  ]
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -30,38 +83,75 @@ def infer(
     randomize_seed,
     width,
     height,
-    model_repo_id=model_repo_id,
+    model_repo_id=model_id_default,
     seed=42,
     guidance_scale=7,
     num_inference_steps=20,
+    model_lora_id=model_lora_default,
+    lora_scale=0.5,
     progress=gr.Progress(track_tqdm=True),
-):
+    ):
+        
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
 
-    pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-    pipe = pipe.to(device)
+    # убираем обновление pipe всегда
+    #pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+    #pipe = pipe.to(device)
+    
+    # добавляем обновление pipe по условию
+    if model_repo_id != model_id_default:
+        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype).to(device)
+        prompt_embeds = process_prompt(prompt, pipe.tokenizer, pipe.text_encoder)
+        negative_prompt_embeds = process_prompt(negative_prompt, pipe.tokenizer, pipe.text_encoder)
+        prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
+    else:
+        # добавляем lora
+        pipe = get_lora_sd_pipeline(ckpt_dir='./lora_lady_and_cats_logos', base_model_name_or_path=model_id_default, dtype=torch_dtype).to(device)
+        #pipe = get_lora_sd_pipeline(ckpt_dir='./'+model_lora_id, base_model_name_or_path=model_id_default, dtype=torch_dtype).to(device)
+        prompt_embeds = process_prompt(prompt, pipe.tokenizer, pipe.text_encoder)
+        negative_prompt_embeds = process_prompt(negative_prompt, pipe.tokenizer, pipe.text_encoder)
+        prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
+        print(f"LoRA adapter loaded: {pipe.unet.active_adapters}")
+        print(f"LoRA scale applied: {lora_scale}")
+        pipe.fuse_lora(lora_scale=lora_scale)
 
-    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]
+
+    # заменяем просто вызов pipe с промптом
+    #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]
+    
+    
+    # на вызов pipe с эмбеддингами
+    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]    
 
     return image, seed
 
 
 examples = [
-    "A young lady in a Russian embroidered kaftan is sitting on a beautiful carved veranda, holding a cup to her mouth and drinking tea from the cup. With her other hand, the girl holds a saucer. The cup and saucer are painted with gzhel. Next to the girl on the table stands a samovar, and steam can be seen above it.",
     "Puss in Boots wearing a sombrero crosses the Grand Canyon on a tightrope with a guitar.",
     "A cat is playing a song called ""About the Cat"" on an accordion by the sea at sunset. The sun is quickly setting behind the horizon, and the light is fading.",
     "A cat walks through the grass on the streets of an abandoned city. The camera view is always focused on the cat's face.",
+    "A young lady in a Russian embroidered kaftan is sitting on a beautiful carved veranda, holding a cup to her mouth and drinking tea from the cup. With her other hand, the girl holds a saucer. The cup and saucer are painted with gzhel. Next to the girl on the table stands a samovar, and steam can be seen above it.",
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
     "An astronaut riding a green horse",
     "A delicious ceviche cheesecake slice",
@@ -158,6 +248,23 @@ with gr.Blocks(css=css) as demo:
                     step=1,
                     value=20,  # Replace with defaults that work for your model
                 )
+                
+                model_lora_id = gr.Dropdown(
+                    label="Lora Id",
+                    choices=model_dropdown,
+                    info="Choose LoRA model",
+                    visible=True,
+                    allow_custom_value=True,
+                    value=model_lora_id,
+                )            
+            
+                lora_scale = gr.Slider(
+                    label="LoRA scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.1,
+                    value=0.5,
+                )                
 
         gr.Examples(examples=examples, inputs=[prompt])
     gr.on(
@@ -173,6 +280,8 @@ with gr.Blocks(css=css) as demo:
             seed,
             guidance_scale,
             num_inference_steps,
+            model_lora_id,
+            lora_scale,
         ],
         outputs=[result, seed],
     )