Update app.py

app.py

@@ -11,6 +11,9 @@ import torch
 from peft import PeftModel, LoraConfig
 import os
 
+# Добавляем глобальный кэш для пайплайнов
+pipe_cache = {}
+
 def get_lora_sd_pipeline(
     ckpt_dir='./lora_logos', 
     base_model_name_or_path=None, 
@@ -33,14 +36,15 @@ def get_lora_sd_pipeline(
     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)))
+    print("UNet 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()
+        if pipe.text_encoder is not None:
+            pipe.text_encoder.half()
     
     return pipe
 
@@ -94,43 +98,50 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
     ):
         
+    global pipe_cache
+ 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
 
-    # убираем обновление 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 = DiffusionPipeline.from_pretrained(model_repo_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)
+    # Кэширование пайплайнов
+    cache_key = f"{model_repo_id}_{model_lora_id}"
+    if cache_key not in pipe_cache:
+        if model_repo_id != model_id_default:
+            pipe = DiffusionPipeline.from_pretrained(model_repo_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:
+            pipe = get_lora_sd_pipeline(
+                ckpt_dir='./'+model_lora_id,
+                base_model_name_or_path=model_id_default,
+                dtype=torch_dtype
+            ).to(device)
+        
+        pipe_cache[cache_key] = pipe
     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}")
+        pipe = pipe_cache[cache_key]
+
+    # Динамическое применение масштаба LoRA
+    if model_repo_id == model_id_default:
+        # Убираем fuse_lora()
+        # pipe.fuse_lora(lora_scale=lora_scale)  # Закомментировали проблемную строку
+        
+        # Вместо этого устанавливаем адаптеры динамически
+        pipe.unet.set_adapters(
+            [model_lora_id], 
+            adapter_weights=[lora_scale]
+        )
+        if hasattr(pipe, 'text_encoder') and pipe.text_encoder is not None:
+            pipe.text_encoder.set_adapters(
+                [model_lora_id], 
+                adapter_weights=[lora_scale]
+            )
+        
+        print(f"Active adapters - UNet: {pipe.unet.active_adapters}, Text Encoder: {pipe.text_encoder.active_adapters if hasattr(pipe, 'text_encoder') else None}")
         print(f"LoRA scale applied: {lora_scale}")
-        pipe.fuse_lora(lora_scale=lora_scale)
-
-
-    # заменяем просто вызов 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 с эмбеддингами