Create utils.py

utils.py

@@ -0,0 +1,49 @@
+import cv2
+import numpy as np
+from PIL import Image
+from transformers import BlipProcessor, BlipForConditionalGeneration
+import torch
+from config import Config
+
+def resize_image_to_1mp(image):
+    """Resizes image to approx 1MP (e.g., 1024x1024) preserving aspect ratio."""
+    w, h = image.size
+    target_pixels = 1024 * 1024
+    aspect_ratio = w / h
+    
+    # Calculate new dimensions
+    new_h = int((target_pixels / aspect_ratio) ** 0.5)
+    new_w = int(new_h * aspect_ratio)
+    
+    # Ensure divisibility by 8 (vae requirement), usually 32 for safety
+    new_w = (new_w // 32) * 32
+    new_h = (new_h // 32) * 32
+    
+    return image.resize((new_w, new_h), Image.LANCZOS)
+
+# Simple caching for captioner
+captioner_processor = None
+captioner_model = None
+
+def get_caption(image):
+    global captioner_processor, captioner_model
+    
+    if captioner_model is None:
+        print("Loading Captioner...")
+        captioner_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+        captioner_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(Config.DEVICE)
+
+    inputs = captioner_processor(image, return_tensors="pt").to(Config.DEVICE)
+    out = captioner_model.generate(**inputs)
+    caption = captioner_processor.decode(out[0], skip_special_tokens=True)
+    return caption
+
+def prepare_control_images(image, zoe_detector, lineart_detector):
+    """Generates the conditioning maps from the input image."""
+    # 1. Zoe Depth Map
+    depth_map = zoe_detector(image, detect_resolution=1024, image_resolution=1024)
+    
+    # 2. LineArt Map
+    lineart_map = lineart_detector(image, detect_resolution=1024, image_resolution=1024)
+    
+    return depth_map, lineart_map