Update app.py

app.py

@@ -4,6 +4,7 @@ import torch
 from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline
 from peft import PeftModel, LoraConfig
 import os
+from PIL import Image
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -65,6 +66,18 @@ pipe_controlnet = StableDiffusionControlNetPipeline.from_pretrained(
     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, 
     negative_prompt, 
@@ -84,13 +97,16 @@ def infer(
     generator = torch.Generator(device).manual_seed(seed)
     
     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
         image = pipe_controlnet(
             prompt=prompt,
             negative_prompt=negative_prompt,
-            image=source_image,     ####################
-            control_image=control_image,  ###############
-            #image=control_image,  # Используем загруженное изображение как карту позы
+            image=source_image,     
+            control_image=control_image,  
             width=width,
             height=height,
             num_inference_steps=num_inference_steps,