Update app.py

app.py

@@ -1,3 +1,5 @@
+# app.py 07.02.25
+
 import gradio as gr
 import numpy as np
 import torch
@@ -17,15 +19,11 @@ def get_lora_sd_pipeline(
     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")
     
-    # Проверка существования директорий LoRA
-    print(f"LoRA directory exists: {os.path.exists(lora_dir)}")
-    print(f"UNet LoRA exists: {os.path.exists(unet_sub_dir)}")
-    print(f"Text encoder LoRA exists: {os.path.exists(text_encoder_sub_dir)}")
-    
     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
@@ -34,30 +32,14 @@ def get_lora_sd_pipeline(
         raise ValueError("Укажите название базовой модели или путь к ней")
     
     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()
     
-    # Логирование параметров до применения LoRA
-    before_params = list(pipe.unet.parameters())
-    
-    # Применение LoRA к UNet
-    if os.path.exists(unet_sub_dir):
-        pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
-        pipe.unet.set_adapter(adapter_name)
-    
-    # Применение LoRA к текстовому энкодеру (если есть)
     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)
     
-    # Логирование параметров после применения LoRA
-    after_params = list(pipe.unet.parameters())
-    print(f"Parameters changed: {before_params != after_params}")
-    
-    # Детальное сравнение параметров
-    for i, (param1, param2) in enumerate(zip(before_params, after_params)):
-        if not torch.equal(param1, param2):
-            print(f"Parameter {i} changed.")
-        else:
-            print(f"Parameter {i} did not change.")
-    
     if dtype in (torch.float16, torch.bfloat16):
         pipe.unet.half()
         pipe.text_encoder.half()
@@ -89,8 +71,8 @@ def infer(
     guidance_scale=7.5, 
     lora_scale=0.5,
     progress=gr.Progress(track_tqdm=True)
-):
-    print(f"Received lora_scale: {lora_scale}")  # Лог для проверки значения lora_scale
+    ):
+    
     generator = torch.Generator(device).manual_seed(seed)
     
     if model != model_default:
@@ -103,13 +85,7 @@ def infer(
         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)
-        
-        # Логирование параметров до и после применения LoRA
-        before_params = list(pipe.unet.parameters())
-        print(f"Applying LoRA with scale: {lora_scale}")
-        pipe.fuse_lora(lora_scale=lora_scale)
-        after_params = list(pipe.unet.parameters())
-        print(f"Parameters changed: {before_params != after_params}")
+        pipe.fuse_lora(lora_scale=lora_scale) # Коэфф. добавления lora
     
     params = {
         'prompt_embeds': prompt_embeds,
@@ -222,6 +198,81 @@ 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,
+                    maximum=1.0,
+                    value=0.5,
+                    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",
+                )
+        
+                # Окно для загрузки изображений
+                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,
+            )
+
         gr.Examples(examples=examples, inputs=[prompt])
         gr.Examples(examples=examples_negative, inputs=[negative_prompt])
         
@@ -248,4 +299,4 @@ with gr.Blocks(css=css) as demo:
 
 if __name__ == "__main__":
     demo.launch()
-        
+