Update app.py

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 import numpy as np
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline, StableDiffusionControlNetImg2ImgPipeline
 from peft import PeftModel, LoraConfig
 import os
+from PIL import Image
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -12,13 +13,15 @@ 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
 
+# Инициализация ControlNet
+controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-openpose", torch_dtype=torch_dtype)
+
 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")
     
@@ -57,6 +60,23 @@ def align_embeddings(prompt_embeds, negative_prompt_embeds):
            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)
+pipe_controlnet = StableDiffusionControlNetPipeline.from_pretrained(
+    model_default,
+    controlnet=controlnet,
+    torch_dtype=torch_dtype
+).to(device)
+
+def preprocess_image(image, target_width, target_height):
+    """
+    Преобразует изображение в формат, подходящий для модели.
+    """
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    image = image.resize((target_width, target_height), Image.LANCZOS)
+    image = np.array(image).astype(np.float32) / 255.0  # Нормализация [0, 1]
+    image = image[None].transpose(0, 3, 1, 2)  # Преобразуем в (batch, channels, height, width)
+    image = torch.from_numpy(image).to(device)
+    return image
 
 def infer(
     prompt, 
@@ -68,57 +88,97 @@ def infer(
     seed=4, 
     guidance_scale=7.5, 
     lora_scale=0.5,
-    # use_control_net=False,  # Добавляем параметр для управления включением ControlNet
-    # control_strength=0.5,  
-    # use_ip_adapter=False,  # Добавляем параметр для управления включением Ip_Adapter
-    # ip_adapter_scale=0.5,
+    use_control_net=False,  # Параметр для включения ControlNet
+    control_strength=0.5,   # Сила влияния ControlNet
+    source_image=None,      # Исходное изображение
+    control_image=None,     # Контрольное изображение
     progress=gr.Progress(track_tqdm=True)
-    ):
-    
+):
     generator = torch.Generator(device).manual_seed(seed)
     
-    if model != model_default:
-        pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device)
+    if use_control_net and control_image is not None and source_image is not None:
+        # Преобразуем изображения
+        source_image = preprocess_image(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)
+        
+        # Используем 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=source_image,     
+            control_image=control_image,  
+            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) # Коэфф. добавления lora
-    
-    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,
-    }
-    
-    # if use_control_net:  # Если ControlNet включен
-    #     params['control_net'] = True  # Включаем использование ControlNet
-    #     params['control_strength'] = control_strength  # Устанавливаем вес ControlNet
-
-    # if use_ip_adapter:  # Если Ip_Adapter включен
-    #     params['ip_adapter'] = True  # Включаем использование Ip_Adapter
-    #     params['ip_adapter_scale'] = ip_adapter_scale  # Устанавливаем вес Ip_Adapter    
+        # Стандартная генерация без ControlNet
+        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,
+        }
+        
+        image = pipe(**params).images[0]
     
-    return pipe(**params).images[0]
+    return image
 
 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.",
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k.",
-    "An astronaut riding a green horse.",
 ]    
 
 examples_negative = [
     "blurred details, low resolution, poor image of a man's face, poor quality, artifacts, black and white image",
-    "blurry details, low resolution, poorly defined edges",
-    "bad face, bad quality, artifacts, low-res, black and white",
 ]
 
 css = """
@@ -130,15 +190,7 @@ css = """
 
 available_models = [
     "stable-diffusion-v1-5/stable-diffusion-v1-5",
-    "SG161222/Realistic_Vision_V3.0_VAE",
     "CompVis/stable-diffusion-v1-4",
-    "stabilityai/sdxl-turbo",
-    "runwayml/stable-diffusion-v1-5",
-    "sd-legacy/stable-diffusion-v1-5",
-    "prompthero/openjourney",
-    "stabilityai/stable-diffusion-3-medium-diffusers",
-    "stabilityai/stable-diffusion-3.5-large",
-    "stabilityai/stable-diffusion-3.5-large-turbo",
 ]
 
 with gr.Blocks(css=css) as demo:
@@ -220,27 +272,14 @@ with gr.Blocks(css=css) as demo:
                     value=512,
                 )
 
-        # Функция для работы с ControlNet ---------------------------------------------------------------------
-        def process_input_ControlNet(image, use_control_net, control_strength, control_mode):
-            if use_control_net:
-                # Логика для обработки с использованием ControlNet
-                result = f"ControlNet активен! Режим: {control_mode}, Интенсивность: {control_strength}"
-            else:
-                # Логика для обработки без ControlNet
-                result = "ControlNet отключен."
-            return result
-
         with gr.Blocks():
             with gr.Row():
-                # Чекбокс для включения/отключения ControlNet
                 use_control_net = gr.Checkbox(
                     label="Use ControlNet",
                     value=False,
                 )
 
-            # Дополнительные опции для ControlNet
             with gr.Column(visible=False) as control_net_options:    
-                # Слайдер для настройки интенсивности
                 control_strength = gr.Slider(
                     label="Control Strength",
                     minimum=0.0,
@@ -249,112 +288,30 @@ with gr.Blocks(css=css) as demo:
                     step=0.05,
                 )
         
-                # Выпадающий список для выбора режима
                 control_mode = gr.Dropdown(
                     label="Control Mode",
                     choices=[
-                        "edge_detection",
-                        "canny_edge_detection",
                         "pose_estimation",
-                        "depth_map",
-                        "segmentation_map",
-                        "scribble_sketch",
-                        "normal_map",
-                        "hed_edge_detection",
-                        "openpose",
-                        "mlsd_line_detection",
-                        "scribble_diffusion",
-                        "semantic_segmentation",
-                        "style_transfer",
-                        "colorization",
-                        "custom_map"
                     ],
                     value="pose_estimation",
                 )
         
-                # Окно для загрузки изображений
+                source_image = gr.Image(label="Upload Source Image")
+                
                 control_image = gr.Image(label="Upload Control Image")
 
-                # Кнопка для запуска работы ControlNet
-                run_button = gr.Button("Run")
-
-                # Текстовое поле для вывода результата
-                output = gr.Textbox(label="Output")
-
-            # Логика для отображения/скрытия дополнительных опций
             use_control_net.change(
                 fn=lambda x: gr.Row.update(visible=x),
                 inputs=use_control_net,
                 outputs=control_net_options,
             )
 
-            # Привязка кнопки Run к функции работы с ControlNet
-            run_button.click(
-                fn=process_input_ControlNet,
-                inputs=[control_image, use_control_net, control_strength, control_mode],
-                outputs=output,
-            )
-
-        # Функция для работы с IP-adapter ----------------------------------------------------------------------------
-        def process_input_IP_adapter(image, use_ip_adapter, ip_adapter_scale, ip_adapter_image):
-            if use_ip_adapter:
-                # Логика для обработки с использованием IP-adapter
-                result = f"IP-adapter активен! Масштаб: {ip_adapter_scale}"
-            else:
-                # Логика для обработки без IP-adapter
-                result = "IP-adapter отключен."
-            return result
-
-        # Создание интерфейса
-        with gr.Blocks():
-            with gr.Row():
-                # Чекбокс для включения/отключения IP-adapter
-                use_ip_adapter = gr.Checkbox(
-                    label="Use IP-adapter",
-                    value=False,
-                )
-
-            # Дополнительные опции для IP-adapter
-            with gr.Column(visible=False) as ip_adapter_options:
-                # Слайдер для настройки масштаба
-                ip_adapter_scale = gr.Slider(
-                     label="IP-adapter Scale",
-                     minimum=0.0,
-                     maximum=1.0,
-                     value=0.5,
-                     step=0.05,
-                )
-        
-                # Окно для загрузки изображений
-                ip_adapter_image = gr.Image(label="Upload IP-adapter Image")
-
-                # Кнопка для запуска обработки
-                run_button = gr.Button("Run")
-
-                # Текстовое поле для вывода результата
-                output = gr.Textbox(label="Output")
-
-            # Логика для отображения/скрытия дополнительных опций
-            use_ip_adapter.change(
-               fn=lambda x: gr.Row.update(visible=x),
-                inputs=use_ip_adapter,
-                outputs=ip_adapter_options,
-            )
-
-            # Привязка кнопки Run к функции работы с IP-adapter
-            run_button.click(
-                fn=process_input_IP_adapter,
-                inputs=[ip_adapter_image, use_ip_adapter, ip_adapter_scale, ip_adapter_image],
-                outputs=output,
-            )
-
         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,
@@ -368,10 +325,10 @@ with gr.Blocks(css=css) as demo:
             seed,
             guidance_scale,
             lora_scale,
-            # use_control_net,  
-            # control_strength,  
-            # use_ip_adapter,  
-            # ip_adapter_scale,
+            use_control_net,   # Добавляем чекбокс для ControlNet
+            control_strength,  # Добавляем контроль силы
+            source_image,      # Добавляем исходное изображение 
+            control_image,     # Добавляем контрольное изображение
         ],
         outputs=[result],
     )