Update app.py

app.py

@@ -11,16 +11,13 @@ 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, 
     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")
     
@@ -33,28 +30,17 @@ def get_lora_sd_pipeline(
     
     pipe = DiffusionPipeline.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)
-    # Исправляем загрузку конфигурации
-    config = LoraConfig.from_pretrained(unet_sub_dir)
-    
-    pipe.unet = PeftModel.from_pretrained(
-        pipe.unet, 
-        unet_sub_dir, 
-        adapter_name=adapter_name,
-        config=config  # Явно передаем конфигурацию
-    )
-    
+    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("UNet Parameters changed:", any(torch.any(b != a) for b, a in zip(before_params, after_params)))
+    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()
-        if pipe.text_encoder is not None:
-            pipe.text_encoder.half()
+        pipe.text_encoder.half()
     
     return pipe
 
@@ -108,51 +94,43 @@ 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)
 
-    # Кэширование пайплайнов
-    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
+    # убираем обновление 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)
     else:
-        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("UNet first layer weights:", pipe.unet.base_model.model[0].weight.data[0,0,:5])
+        # добавляем 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)
+
+
+    # заменяем просто вызов 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 с эмбеддингами