Update app.py

app.py

@@ -94,7 +94,8 @@ def infer(
     progress=gr.Progress(track_tqdm=True)
 ):
     generator = torch.Generator(device).manual_seed(seed)
-    
+
+    # Генерация с IP_adapter
     if use_control_net and control_image is not None and cn_source_image is not None:
         # pipe_controlnet = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(    
         #     model_default,
@@ -153,30 +154,89 @@ def infer(
             generator=generator
         ).images[0]
     else:
-        # Стандартная генерация без ControlNet
-        if model != model_default:
-            pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device)
+        # Генерация с ControlNet
+        if use_control_net and control_image is not None and cn_source_image is not None:
+            # pipe_controlnet = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(    
+            #     model_default,
+            #     controlnet=controlnet,
+            #     torch_dtype=torch_dtype
+            # ).to(device)
+        
+            # Преобразуем изображения
+            cn_source_image = preprocess_image(cn_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)
+            #strength_sn = float(strength_sn)
+        
+            # Используем 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=cn_source_image,     
+                control_image=control_image,  
+                strength=strength_cn, # Коэфф. зашумления, чем больше, тем больше меняется результирующее изображение относитенльно исходного
+                width=width,
+                height=height,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                controlnet_conditioning_scale=control_strength,
+                generator=generator
+            ).images[0]
         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)
+            # Генерация без ControlNet и IP_adapter
+            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,
-        }
+            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]
+            image = pipe(**params).images[0]
     
     return image