first commit

.vscode/launch.json

@@ -0,0 +1,77 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "train mode a (debug)",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/src/train_cn_sdxl_mc_v2.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "env": {
+                "PYTORCH_ALLOC_CONF": "expandable_segments:True",
+                "XFORMERS_DISABLED": "1"
+            },
+            "args": [
+                "--data_root", "/mnt/data1tb/CarlaDS/all_collected_ds",
+                "--output_dir", "./checkpoints_debug",
+                "--mode", "a",
+                "--mixed_precision", "fp16",
+                "--image_size", "512",
+                "--train_batch_size", "1",
+                "--grad_accum", "1",
+                "--max_steps", "10",
+                "--save_every", "5",
+                "--num_workers", "0"
+            ]
+        },
+        {
+            "name": "train mode b1 (debug)",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/src/train_cn_sdxl_mc_v2.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "env": {
+                "PYTORCH_ALLOC_CONF": "expandable_segments:True",
+                "XFORMERS_DISABLED": "1"
+            },
+            "args": [
+                "--data_root", "/mnt/data1tb/CarlaDS/all_collected_ds",
+                "--output_dir", "./checkpoints_debug",
+                "--mode", "b1",
+                "--mixed_precision", "fp16",
+                "--image_size", "512",
+                "--train_batch_size", "1",
+                "--grad_accum", "1",
+                "--max_steps", "5",
+                "--save_every", "5",
+                "--num_workers", "0"
+            ]
+        },
+        {
+            "name": "train full (mode a)",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/src/train_cn_sdxl_mc_v2.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "env": {
+                "PYTORCH_ALLOC_CONF": "expandable_segments:True",
+                "XFORMERS_DISABLED": "1"
+            },
+            "args": [
+                "--data_root", "/mnt/data1tb/CarlaDS/all_collected_ds",
+                "--output_dir", "./checkpoints",
+                "--mode", "a",
+                "--mixed_precision", "fp16",
+                "--image_size", "512",
+                "--train_batch_size", "1",
+                "--grad_accum", "8",
+                "--max_steps", "50000",
+                "--save_every", "2000",
+                "--num_workers", "4"
+            ]
+        }
+    ]
+}

.vscode/settings.json

@@ -0,0 +1,4 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:conda"
+}

configs/depth2street_dataset.gin

@@ -0,0 +1,28 @@
+# Запуск:
+# python -m src.run_depth2street --gin configs/depth2street_dataset.gin
+
+Depth2StreetGenerator.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+Depth2StreetGenerator.controlnet_id = "diffusers/controlnet-depth-sdxl-1.0"
+
+Depth2StreetGenerator.device = "cuda"
+Depth2StreetGenerator.torch_dtype = "fp16"
+Depth2StreetGenerator.enable_xformers = False
+
+Depth2StreetGenerator.size = 1024
+Depth2StreetGenerator.steps = 30
+Depth2StreetGenerator.guidance_scale = 5.0
+Depth2StreetGenerator.control_scale = 1.0
+Depth2StreetGenerator.seed = 123
+
+Depth2StreetGenerator.prompt = "realistic street view photo, eye-level camera, wide angle lens, urban street, buildings, road, natural lighting, high detail"
+Depth2StreetGenerator.negative_prompt = "aerial view, top-down, bird's-eye view, map, satellite, isometric, low quality, blurry, distortion, warped perspective"
+
+# dataset mode
+Depth2StreetGenerator.ds_root = "ds"
+Depth2StreetGenerator.depth_filename = "depth.npy"
+Depth2StreetGenerator.out_dir = "out_depth2street"
+Depth2StreetGenerator.out_name = "street_gen.png"
+Depth2StreetGenerator.overwrite = False
+
+# single mode не используется
+Depth2StreetGenerator.depth_npy_path = None

configs/depth2street_dronly_masks.gin

@@ -0,0 +1,36 @@
+DepthEdgeDepth2StreetGenerator.ds_root = "/media/servml/SSD_1/ds"
+
+# файлы в каждой папке семпла:
+DepthEdgeDepth2StreetGenerator.edge_filename = "depth_drone_edge.png"
+DepthEdgeDepth2StreetGenerator.depth_filename = "depth_drone.png"
+
+DepthEdgeDepth2StreetGenerator.out_dir = "out_street_expC"
+DepthEdgeDepth2StreetGenerator.out_name = "street_gen.png"
+DepthEdgeDepth2StreetGenerator.overwrite = False
+
+# модели
+DepthEdgeDepth2StreetGenerator.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+# используем depth-controlnet и для edge, и для depth (works ok на старте)
+DepthEdgeDepth2StreetGenerator.controlnet_edge_id = "diffusers/controlnet-depth-sdxl-1.0"
+DepthEdgeDepth2StreetGenerator.controlnet_depth_id = "diffusers/controlnet-depth-sdxl-1.0"
+
+DepthEdgeDepth2StreetGenerator.device = "cuda"
+DepthEdgeDepth2StreetGenerator.torch_dtype = "fp16"
+DepthEdgeDepth2StreetGenerator.enable_xformers = False
+
+# генерация
+DepthEdgeDepth2StreetGenerator.size = 1024
+DepthEdgeDepth2StreetGenerator.steps = 30
+DepthEdgeDepth2StreetGenerator.guidance_scale = 5.0
+DepthEdgeDepth2StreetGenerator.seed = 123
+
+# ключевое для Exp C:
+DepthEdgeDepth2StreetGenerator.edge_scale = 1.0
+DepthEdgeDepth2StreetGenerator.depth_scale = 0.5
+
+DepthEdgeDepth2StreetGenerator.prompt = "realistic street-level photo, eye-level viewpoint, camera height 1.6m, wide-angle 24mm lens, street, sidewalks, buildings facades, road, natural daylight, high detail, sharp focus"
+DepthEdgeDepth2StreetGenerator.negative_prompt = "aerial view, drone view, top-down, bird's-eye view, satellite, map, isometric, miniature, tilt-shift, orthographic, lowres, blurry, warped perspective, distorted geometry"
+
+# если хочешь прогнать первые N сэмплов:
+DepthEdgeDepth2StreetGenerator.max_samples = 0
+

configs/depth2street_single.gin

@@ -0,0 +1,24 @@
+# Запуск:
+# python -m src.run_depth2street --gin configs/depth2street_single.gin
+
+Depth2StreetGenerator.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+Depth2StreetGenerator.controlnet_id = "diffusers/controlnet-depth-sdxl-1.0"
+
+Depth2StreetGenerator.device = "cuda"
+Depth2StreetGenerator.torch_dtype = "fp16"
+Depth2StreetGenerator.enable_xformers = False
+
+Depth2StreetGenerator.size = 1024
+Depth2StreetGenerator.steps = 30
+Depth2StreetGenerator.guidance_scale = 5.0
+Depth2StreetGenerator.control_scale = 1.0
+Depth2StreetGenerator.seed = 123
+
+Depth2StreetGenerator.prompt = "realistic street view photo, eye-level camera, wide angle lens, urban street, buildings, road, natural lighting, high detail"
+Depth2StreetGenerator.negative_prompt = "aerial view, top-down, bird's-eye view, map, satellite, isometric, low quality, blurry, distortion, warped perspective"
+
+Depth2StreetGenerator.depth_npy_path = "ds/00001/depth.npy"
+Depth2StreetGenerator.out_path = "out_single.png"
+
+# dataset mode не используется
+Depth2StreetGenerator.ds_root = None

configs/depth_drone.gin

@@ -0,0 +1,14 @@
+# DepthGenerator.ds_root = "/media/servml/SSD_1/ds"
+DepthGenerator.ds_root = "/media/servml/SSD_1/21"
+DepthGenerator.source = "drone"
+
+DepthGenerator.model_id = "depth-anything/Depth-Anything-V2-Small-hf"
+DepthGenerator.device = "cuda"
+
+DepthGenerator.out_name = "depth.npy"
+DepthGenerator.out_vis_name = "depth_vis.png"
+
+DepthGenerator.exts = ["png", "jpg", "jpeg"]
+DepthGenerator.recursive = False
+DepthGenerator.overwrite = False
+DepthGenerator.batch_log_every = 50

configs/depth_edge_gen.gin

@@ -0,0 +1,17 @@
+# DepthEdgesGenerator.ds_root = "/media/servml/SSD_1/ds"
+DepthEdgesGenerator.ds_root = "/media/servml/SSD_1/21"
+
+DepthEdgesGenerator.depth_filename = "depth_drone.npy"
+
+DepthEdgesGenerator.out_name = "depth_drone_edge.png"
+
+DepthEdgesGenerator.p_lo = 2.0
+DepthEdgesGenerator.p_hi = 98.0
+DepthEdgesGenerator.blur_ksize = 5
+
+DepthEdgesGenerator.recursive = False
+DepthEdgesGenerator.overwrite = False
+DepthEdgesGenerator.batch_log_every = 50
+
+
+

configs/depth_sat.gin

@@ -0,0 +1,13 @@
+DepthGenerator.ds_root = "/media/servml/SSD_1/ds"
+DepthGenerator.source = "sat"
+
+DepthGenerator.model_id = "depth-anything/Depth-Anything-V2-Small-hf"
+DepthGenerator.device = "cuda"
+
+DepthGenerator.out_name = "depth.npy"
+DepthGenerator.out_vis_name = "depth_vis.png"
+
+DepthGenerator.exts = ["png", "jpg", "jpeg"]
+DepthGenerator.recursive = False
+DepthGenerator.overwrite = False
+DepthGenerator.batch_log_every = 50

notes.txt

@@ -0,0 +1,93 @@
+PYTORCH_ALLOC_CONF=expandable_segments:True \
+XFORMERS_DISABLED=1 \
+~/miniconda3/envs/streetview/bin/python /home/servml/Документы/code/Bogdan/generator/src/train_cn_sdxl_mc_v2.py \
+    --data_root /mnt/data1tb/CarlaDS/all_collected_ds \
+    --output_dir ./checkpoints_safe \
+    --mode a \
+    --mixed_precision fp16 \
+    --image_size 512 \
+    --train_batch_size 1 \
+    --grad_accum 8 \
+    --max_steps 2000 \
+    --save_every 500
+
+
+
+PYTORCH_ALLOC_CONF=expandable_segments:True \
+XFORMERS_DISABLED=1 \
+~/miniconda3/envs/streetview/bin/python \
+    /home/servml/Документы/code/Bogdan/generator/src/train_cn_sdxl_mc_v2.py \
+    --data_root /mnt/data1tb/CarlaDS/all_collected_ds \
+    --output_dir ./checkpoints \
+    --mode a \
+    --mixed_precision fp16 \
+    --image_size 512 \
+    --train_batch_size 2 \
+    --grad_accum 8 \
+    --max_steps 5 \
+    --save_every 5
+
+------------------------------------------------------------------------------------------
+
+Добавляю функцию sanity_sample и вызов в цикле. Изменения минимальные — только то что нужно.
+1. Добавьте импорты в начало файла:
+
+from diffusers import StableDiffusionXLControlNetPipeline
+import numpy as np
+
+2. Добавьте функцию перед main():
+
+@torch.no_grad()
+def run_sanity_sample(accelerator, controlnet, args, batch, step, device):
+    """Генерит 1 картинку на фиксированном seed и сохраняет в output_dir/samples/."""
+    ensure_dir(os.path.join(args.output_dir, "samples"))
+
+    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+        args.pretrained_model,
+        controlnet=accelerator.unwrap_model(controlnet),
+        torch_dtype=torch.float16,
+        variant="fp16",
+    ).to(device)
+    pipe.set_progress_bar_config(disable=True)
+
+    # Берём первый сэмпл из батча как control
+    control_img = batch["conditioning_pixel_values"][0].cpu()          # [C, H, W], float16 [-1,1] или [0,1]
+    control_img = (control_img.float().clamp(-1, 1) + 1) / 2          # → [0,1]
+    control_img = (control_img.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
+    control_pil = Image.fromarray(control_img)
+
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    image = pipe(
+        prompt=args.prompt_default,
+        image=control_pil,
+        num_inference_steps=20,
+        guidance_scale=7.5,
+        generator=generator,
+        height=args.image_size,
+        width=args.image_size,
+    ).images[0]
+
+    out_path = os.path.join(args.output_dir, "samples", f"step_{step:06d}.png")
+    image.save(out_path)
+    print(f"\n[SAMPLE] saved → {out_path}")
+
+    del pipe
+    torch.cuda.empty_cache()
+
+3. В цикле обучения — замените блок _save на:
+
+if accelerator.is_main_process and global_step % args.save_every == 0:
+    _save(accelerator, controlnet, args.output_dir, global_step)
+    run_sanity_sample(accelerator, controlnet, args, batch, global_step, device)
+
+    Что делает:
+
+Берёт первый control из текущего батча (фиксированный контент)
+Генерит на фиксированном --seed → результат воспроизводим
+Сохраняет в output_dir/samples/step_000500.png, step_001000.png и т.д.
+Удаляет pipeline после генерации → не держит лишний VRAM
+
+Если control-каналов несколько (mode b/c), [0] возьмёт первые 3 канала — для визуального контроля этого достаточно.
+
+

src/dataset.py

@@ -0,0 +1,287 @@
+"""
+StreetViewControlDataset
+========================
+Датасет для обучения ControlNet SDXL на данных CARLA.
+
+Структура датасета:
+    all_collected_ds/
+    ├── cell_1/
+    │   ├── satellite/
+    │   │   ├── rgb.png
+    │   │   ├── semantic_color.png
+    │   │   ├── depth_uint16.png
+    │   │   └── ...
+    │   ├── drone/
+    │   │   ├── yaw180_rgb.png
+    │   │   ├── yaw180_depth_uint16.png
+    │   │   └── ...
+    │   └── street/
+    │       ├── forward_rgb.png    ← TARGET (ground truth)
+    │       └── ...
+    └── cell_2/
+        └── ...
+
+Режимы (mode):
+    "a" — controls: [sat/semantic_color]
+    "b" — controls: [sat/semantic_color, drone/rgb]
+    "c" — controls: [sat/semantic_color, sat/depth,
+                     drone/rgb,          drone/depth]
+
+Все режимы: target = street/forward_rgb (ground truth).
+Street view НЕ используется как input — только как обучающая цель.
+Depth спутника и дрона даёт геометрию сцены сверху.
+
+Возвращает dict:
+    pixel_values              : [3,H,W] float32 [-1,1]   (target)
+    conditioning_pixel_values : list[[3,H,W] float32 [0,1]]  (controls)
+    prompts                   : str
+    cell                      : str  (имя ячейки)
+"""
+
+from pathlib import Path
+from typing import List
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+import torchvision.transforms.functional as TF
+
+
+# ---------------------------------------------------------------------------
+# Image loaders
+# ---------------------------------------------------------------------------
+
+def _load_rgb_target(path: str, size: int) -> torch.Tensor:
+    """RGB → [3,H,W] float32 в [-1, 1]  (для target/pixel_values)."""
+    img = Image.open(path).convert("RGB")
+    img = img.resize((size, size), resample=Image.BILINEAR)
+    t = TF.to_tensor(img)      # [0,1]
+    return t * 2.0 - 1.0      # [-1,1]
+
+
+def _load_rgb_control(path: str, size: int) -> torch.Tensor:
+    """RGB → [3,H,W] float32 в [0, 1]  (для control images)."""
+    img = Image.open(path).convert("RGB")
+    img = img.resize((size, size), resample=Image.BILINEAR)
+    return TF.to_tensor(img)
+
+
+def _load_depth_uint16(path: str, size: int) -> torch.Tensor:
+    """
+    depth_uint16.png → [3,H,W] float32 в [0,1].
+    Нормализация по 2–98 percentile чтобы убрать выбросы.
+    Канал повторяется трижды (grayscale → RGB) для ControlNet.
+    """
+    arr = np.array(Image.open(path), dtype=np.float32)   # HxW
+    lo  = np.percentile(arr, 2.0)
+    hi  = np.percentile(arr, 98.0)
+    arr = np.clip((arr - lo) / (hi - lo + 1e-6), 0.0, 1.0)
+    pil = Image.fromarray((arr * 255).astype(np.uint8), mode="L").convert("RGB")
+    pil = pil.resize((size, size), resample=Image.BILINEAR)
+    return TF.to_tensor(pil)
+
+
+def _blend(rgb: torch.Tensor, seg: torch.Tensor, alpha: float) -> torch.Tensor:
+    """Смешиваем satellite rgb и semantic_color. alpha=1 → только seg."""
+    return (1.0 - alpha) * rgb + alpha * seg
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+class StreetViewControlDataset(Dataset):
+    """
+    Args:
+        data_root         : путь к корню датасета
+        mode              : "a" | "b" | "c"
+        image_size        : разрешение (одинаковое для всех изображений)
+        sat_seg_blend     : alpha смешивания sat/rgb + sat/semantic_color
+                            0.0 = только rgb спутника
+                            1.0 = только semantic карта
+                            0.5 = 50/50 (дефолт)
+        drone_yaw         : префикс yaw для drone (например "yaw180")
+        street_direction  : направление для street target (например "forward")
+        prompt_default    : дефолтный промпт если нет prompt.txt в ячейке
+        augment           : горизонтальный flip с вероятностью 0.5
+    """
+
+    # Файлы необходимые для каждого mode
+    _REQUIRED = {
+        "a": [
+            ("satellite", "semantic_color.png"),
+            ("street",    "{d}_rgb.png"),
+        ],
+        "b": [
+            ("satellite", "semantic_color.png"),
+            ("drone",     "{y}_rgb.png"),
+            ("street",    "{d}_rgb.png"),
+        ],
+        "c": [
+            ("satellite", "semantic_color.png"),
+            ("satellite", "depth_uint16.png"),
+            ("drone",     "{y}_rgb.png"),
+            ("drone",     "{y}_depth_uint16.png"),
+            ("street",    "{d}_rgb.png"),
+        ],
+    }
+
+    def __init__(
+        self,
+        data_root: str,
+        mode: str = "a",
+        image_size: int = 1024,
+        sat_seg_blend: float = 0.5,
+        drone_yaw: str = "yaw180",
+        street_direction: str = "forward",
+        prompt_default: str = (
+            "realistic street view photo, urban road, "
+            "natural lighting, high detail"
+        ),
+        augment: bool = False,
+    ):
+        assert mode in ("a", "b", "c"), \
+            f"mode должен быть 'a', 'b' или 'c', получено: {mode!r}"
+
+        self.root           = Path(data_root)
+        self.mode           = mode
+        self.sz             = image_size
+        self.sat_seg_blend  = sat_seg_blend
+        self.yaw            = drone_yaw
+        self.direction      = street_direction
+        self.prompt_default = prompt_default
+        self.augment        = augment
+
+        self.cells = self._scan()
+        if not self.cells:
+            raise RuntimeError(
+                f"Не найдено валидных ячеек в {data_root} для mode={mode}")
+
+        print(f"[Dataset] mode={mode} | {len(self.cells)} ячеек | "
+              f"size={image_size} | controls={self._num_controls()}")
+
+    def _num_controls(self) -> int:
+        return {"a": 1, "b": 2, "c": 4}[self.mode]
+
+    def _required_files(self) -> List[tuple]:
+        return [
+            (folder, fname.format(d=self.direction, y=self.yaw))
+            for folder, fname in self._REQUIRED[self.mode]
+        ]
+
+    def _is_valid(self, cell: Path) -> bool:
+        for folder, fname in self._required_files():
+            if not (cell / folder / fname).exists():
+                return False
+        return True
+
+    def _scan(self) -> List[Path]:
+        return sorted(
+            c for c in self.root.iterdir()
+            if c.is_dir() and self._is_valid(c)
+        )
+
+    def _prompt(self, cell: Path) -> str:
+        p = cell / "prompt.txt"
+        return p.read_text().strip() if p.exists() else self.prompt_default
+
+    def __len__(self) -> int:
+        return len(self.cells)
+
+    def __getitem__(self, idx: int) -> dict:
+        cell = self.cells[idx]
+        d, y, sz = self.direction, self.yaw, self.sz
+
+        # ── Target ────────────────────────────────────────────────────────────
+        target = _load_rgb_target(
+            str(cell / "street" / f"{d}_rgb.png"), sz)
+
+        # ── Controls ──────────────────────────────────────────────────────────
+        controls: List[torch.Tensor] = []
+
+        # Control: satellite semantic_color (всегда, во всех режимах)
+        sat_seg = _load_rgb_control(
+            str(cell / "satellite" / "semantic_color.png"), sz)
+
+        # Опционально смешиваем с satellite rgb для лучшей читаемости
+        sat_rgb_path = cell / "satellite" / "rgb.png"
+        if self.sat_seg_blend < 1.0 and sat_rgb_path.exists():
+            sat_rgb = _load_rgb_control(str(sat_rgb_path), sz)
+            sat_seg = _blend(sat_rgb, sat_seg, self.sat_seg_blend)
+
+        controls.append(sat_seg)  # control[0] всегда
+
+        # mode c: satellite depth → control[1]
+        if self.mode == "c":
+            controls.append(_load_depth_uint16(
+                str(cell / "satellite" / "depth_uint16.png"), sz))
+
+        # mode b, c: drone rgb
+        if self.mode in ("b", "c"):
+            controls.append(_load_rgb_control(
+                str(cell / "drone" / f"{y}_rgb.png"), sz))
+
+        # mode c: drone depth → последний control
+        if self.mode == "c":
+            controls.append(_load_depth_uint16(
+                str(cell / "drone" / f"{y}_depth_uint16.png"), sz))
+
+        # ── Augmentation ──────────────────────────────────────────────────────
+        if self.augment and torch.rand(1).item() > 0.5:
+            target   = TF.hflip(target)
+            controls = [TF.hflip(c) for c in controls]
+
+        # Конкатенируем все controls по каналам -> [3*N, H, W]
+        # mode a: [3,H,W], mode b: [6,H,W], mode c: [12,H,W]
+        control_concat = torch.cat(controls, dim=0)
+
+        return {
+            "pixel_values":              target,           # [3,H,W]
+            "conditioning_pixel_values": control_concat,  # [3*N,H,W]
+            "num_controls":              len(controls),
+            "prompts":                   self._prompt(cell),
+            "cell":                      cell.name,
+        }
+
+
+# ---------------------------------------------------------------------------
+# Collate
+# ---------------------------------------------------------------------------
+
+def collate_fn(batch: List[dict]) -> dict:
+    pixel_values = torch.stack([b["pixel_values"] for b in batch])
+
+    # conditioning_pixel_values уже конкатенирован в dataset -> [3*N, H, W]
+    conditioning = torch.stack([b["conditioning_pixel_values"] for b in batch])
+
+    return {
+        "pixel_values":              pixel_values,   # [B, 3, H, W]
+        "conditioning_pixel_values": conditioning,   # [B, 3*N, H, W]
+        "num_controls":              batch[0]["num_controls"],
+        "prompts":                   [b["prompts"] for b in batch],
+        "cells":                     [b["cell"]   for b in batch],
+    }
+
+
+# ---------------------------------------------------------------------------
+# Quick test
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import sys
+    root = sys.argv[1] if len(sys.argv) > 1 else "./all_collected_ds"
+    mode = sys.argv[2] if len(sys.argv) > 2 else "a"
+
+    ds = StreetViewControlDataset(root, mode=mode, image_size=512, augment=True)
+    print(f"Размер датасета: {len(ds)}")
+
+    s = ds[0]
+    print(f"pixel_values:  {s['pixel_values'].shape}  "
+          f"[{s['pixel_values'].min():.2f}, {s['pixel_values'].max():.2f}]")
+    for i, c in enumerate(s["conditioning_pixel_values"]):
+        print(f"control[{i}]:   {c.shape}  [{c.min():.2f}, {c.max():.2f}]")
+    print(f"prompt: {s['prompts']}")
+    print(f"cell:   {s['cell']}")
+
+    

src/dataset_v2.py

@@ -0,0 +1,302 @@
+"""
+StreetViewControlDataset
+========================
+Датасет для обучения ControlNet SDXL на данных CARLA.
+
+Структура датасета:
+    all_collected_ds/
+    ├── cell_1/
+    │   ├── satellite/
+    │   │   ├── rgb.png
+    │   │   ├── semantic_color.png
+    │   │   ├── depth_uint16.png
+    │   │   └── ...
+    │   ├── drone/
+    │   │   ├── yaw180_rgb.png
+    │   │   ├── yaw180_depth_uint16.png
+    │   │   └── ...
+    │   └── street/
+    │       ├── forward_rgb.png    ← TARGET (ground truth)
+    │       └── ...
+    └── cell_2/
+        └── ...
+
+Режимы (mode):
+    "a"  — controls: [drone/depth]
+    "b1" — controls: [drone/depth, sat/semantic_color]
+    "b2" — controls: [drone/depth, sat/depth]
+    "c"  — controls: [sat/semantic_color, drone/rgb, street/depth]
+
+Все режимы: target = street/forward_rgb (ground truth).
+Street view НЕ используется как input — только как обучающая цель.
+
+Возвращает dict:
+    pixel_values              : [3,H,W] float32 [-1,1]   (target)
+    conditioning_pixel_values : list[[3,H,W] float32 [0,1]]  (controls)
+    prompts                   : str
+    cell                      : str  (имя ячейки)
+"""
+
+from pathlib import Path
+from typing import List
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+import torchvision.transforms.functional as TF
+
+
+# ---------------------------------------------------------------------------
+# Image loaders
+# ---------------------------------------------------------------------------
+
+def _load_rgb_target(path: str, size: int) -> torch.Tensor:
+    """RGB → [3,H,W] float32 в [-1, 1]  (для target/pixel_values)."""
+    img = Image.open(path).convert("RGB")
+    img = img.resize((size, size), resample=Image.BILINEAR)
+    t = TF.to_tensor(img)      # [0,1]
+    return t * 2.0 - 1.0      # [-1,1]
+
+
+def _load_rgb_control(path: str, size: int) -> torch.Tensor:
+    """RGB → [3,H,W] float32 в [0, 1]  (для control images)."""
+    img = Image.open(path).convert("RGB")
+    img = img.resize((size, size), resample=Image.BILINEAR)
+    return TF.to_tensor(img)
+
+
+def _load_depth_uint16(path: str, size: int) -> torch.Tensor:
+    """
+    depth_uint16.png → [3,H,W] float32 в [0,1].
+    Нормализация по 2–98 percentile чтобы убрать выбросы.
+    Канал повторяется трижды (grayscale → RGB) для ControlNet.
+    """
+    arr = np.array(Image.open(path), dtype=np.float32)   # HxW
+    lo  = np.percentile(arr, 2.0)
+    hi  = np.percentile(arr, 98.0)
+    arr = np.clip((arr - lo) / (hi - lo + 1e-6), 0.0, 1.0)
+    pil = Image.fromarray((arr * 255).astype(np.uint8), mode="L").convert("RGB")
+    pil = pil.resize((size, size), resample=Image.BILINEAR)
+    return TF.to_tensor(pil)
+
+
+def _blend(rgb: torch.Tensor, seg: torch.Tensor, alpha: float) -> torch.Tensor:
+    """Смешиваем satellite rgb и semantic_color. alpha=1 → только seg."""
+    return (1.0 - alpha) * rgb + alpha * seg
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+class StreetViewControlDataset(Dataset):
+    """
+    Args:
+        data_root         : путь к корню датасета
+        mode              : "a" | "b" | "c"
+        image_size        : разрешение (одинаковое для всех изображений)
+        sat_seg_blend     : alpha смешивания sat/rgb + sat/semantic_color
+                            0.0 = только rgb спутника
+                            1.0 = только semantic карта
+                            0.5 = 50/50 (дефолт)
+        drone_yaw         : префикс yaw для drone (например "yaw180")
+        street_direction  : направление для street target (например "forward")
+        prompt_default    : дефолтный промпт если нет prompt.txt в ячейке
+        augment           : горизонтальный flip с вероятностью 0.5
+    """
+
+    # Файлы необходимые для каждого mode
+    _REQUIRED = {
+        "a": [
+            ("drone",   "{y}_depth_uint16.png"),
+            ("street",  "{d}_rgb.png"),
+        ],
+        "b1": [
+            ("drone",     "{y}_depth_uint16.png"),
+            ("satellite", "semantic_color.png"),
+            ("street",    "{d}_rgb.png"),
+        ],
+        "b2": [
+            ("drone",     "{y}_depth_uint16.png"),
+            ("satellite", "depth_uint16.png"),
+            ("street",    "{d}_rgb.png"),
+        ],
+        "c": [
+            ("satellite", "semantic_color.png"),
+            ("drone",     "{y}_rgb.png"),
+            ("street",    "{d}_depth_uint16.png"),
+            ("street",    "{d}_rgb.png"),
+        ],
+    }
+
+    def __init__(
+        self,
+        data_root: str,
+        mode: str = "a",
+        image_size: int = 1024,
+        sat_seg_blend: float = 0.5,
+        drone_yaw: str = "yaw180",
+        street_direction: str = "forward",
+        prompt_default: str = (
+            "realistic street view photo, urban road, "
+            "natural lighting, high detail"
+        ),
+        augment: bool = False,
+    ):
+        print(mode)
+
+        assert mode in ("a", "b1", "b2", "c"), \
+            f"mode должен быть 'a', 'b1', 'b2' или 'c', получено: {mode!r}"
+
+        self.root           = Path(data_root)
+        self.mode           = mode
+        self.sz             = image_size
+        self.sat_seg_blend  = sat_seg_blend
+        self.yaw            = drone_yaw
+        self.direction      = street_direction
+        self.prompt_default = prompt_default
+        self.augment        = augment
+
+        self.cells = self._scan()
+        if not self.cells:
+            raise RuntimeError(
+                f"Не найдено валидных ячеек в {data_root} для mode={mode}")
+
+        print(f"[Dataset] mode={mode} | {len(self.cells)} ячеек | "
+              f"size={image_size} | controls={self._num_controls()}")
+
+    def _num_controls(self) -> int:
+        return {"a": 1, "b1": 2, "b2": 2, "c": 3}[self.mode]
+
+    def _required_files(self) -> List[tuple]:
+        return [
+            (folder, fname.format(d=self.direction, y=self.yaw))
+            for folder, fname in self._REQUIRED[self.mode]
+        ]
+
+    def _is_valid(self, cell: Path) -> bool:
+        for folder, fname in self._required_files():
+            if not (cell / folder / fname).exists():
+                return False
+        return True
+
+    def _scan(self) -> List[Path]:
+        return sorted(
+            c for c in self.root.iterdir()
+            if c.is_dir() and self._is_valid(c)
+        )
+
+    def _prompt(self, cell: Path) -> str:
+        p = cell / "prompt.txt"
+        return p.read_text().strip() if p.exists() else self.prompt_default
+
+    def __len__(self) -> int:
+        return len(self.cells)
+
+    def __getitem__(self, idx: int) -> dict:
+        cell = self.cells[idx]
+        d, y, sz = self.direction, self.yaw, self.sz
+
+        # ── Target ────────────────────────────────────────────────────────────
+        target = _load_rgb_target(
+            str(cell / "street" / f"{d}_rgb.png"), sz)
+
+        # ── Controls ──────────────────────────────────────────────────────────
+        controls: List[torch.Tensor] = []
+
+        if self.mode == "a":
+            # [drone/depth]
+            controls.append(_load_depth_uint16(
+                str(cell / "drone" / f"{y}_depth_uint16.png"), sz))
+
+        elif self.mode == "b1":
+            # [drone/depth, sat/semantic_color]
+            controls.append(_load_depth_uint16(
+                str(cell / "drone" / f"{y}_depth_uint16.png"), sz))
+            sat_seg = _load_rgb_control(
+                str(cell / "satellite" / "semantic_color.png"), sz)
+            sat_rgb_path = cell / "satellite" / "rgb.png"
+            if self.sat_seg_blend < 1.0 and sat_rgb_path.exists():
+                sat_rgb = _load_rgb_control(str(sat_rgb_path), sz)
+                sat_seg = _blend(sat_rgb, sat_seg, self.sat_seg_blend)
+            controls.append(sat_seg)
+
+        elif self.mode == "b2":
+            # [drone/depth, sat/depth]
+            controls.append(_load_depth_uint16(
+                str(cell / "drone" / f"{y}_depth_uint16.png"), sz))
+            controls.append(_load_depth_uint16(
+                str(cell / "satellite" / "depth_uint16.png"), sz))
+
+        elif self.mode == "c":
+            # [sat/semantic_color, drone/rgb, street/depth]
+            sat_seg = _load_rgb_control(
+                str(cell / "satellite" / "semantic_color.png"), sz)
+            sat_rgb_path = cell / "satellite" / "rgb.png"
+            if self.sat_seg_blend < 1.0 and sat_rgb_path.exists():
+                sat_rgb = _load_rgb_control(str(sat_rgb_path), sz)
+                sat_seg = _blend(sat_rgb, sat_seg, self.sat_seg_blend)
+            controls.append(sat_seg)
+            controls.append(_load_rgb_control(
+                str(cell / "drone" / f"{y}_rgb.png"), sz))
+            controls.append(_load_depth_uint16(
+                str(cell / "street" / f"{d}_depth_uint16.png"), sz))
+
+        # ── Augmentation ──────────────────────────────────────────────────────
+        if self.augment and torch.rand(1).item() > 0.5:
+            target   = TF.hflip(target)
+            controls = [TF.hflip(c) for c in controls]
+
+        # Конкатенируем все controls по каналам -> [3*N, H, W]
+        # mode a: [3,H,W], mode b: [6,H,W], mode c: [12,H,W]
+        control_concat = torch.cat(controls, dim=0)
+
+        return {
+            "pixel_values":              target,           # [3,H,W]
+            "conditioning_pixel_values": control_concat,  # [3*N,H,W]
+            "num_controls":              len(controls),
+            "prompts":                   self._prompt(cell),
+            "cell":                      cell.name,
+        }
+
+
+# ---------------------------------------------------------------------------
+# Collate
+# ---------------------------------------------------------------------------
+
+def collate_fn(batch: List[dict]) -> dict:
+    pixel_values = torch.stack([b["pixel_values"] for b in batch])
+
+    # conditioning_pixel_values уже конкатенирован в dataset -> [3*N, H, W]
+    conditioning = torch.stack([b["conditioning_pixel_values"] for b in batch])
+
+    return {
+        "pixel_values":              pixel_values,   # [B, 3, H, W]
+        "conditioning_pixel_values": conditioning,   # [B, 3*N, H, W]
+        "num_controls":              batch[0]["num_controls"],
+        "prompts":                   [b["prompts"] for b in batch],
+        "cells":                     [b["cell"]   for b in batch],
+    }
+
+
+# ---------------------------------------------------------------------------
+# Quick test
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import sys
+    root = sys.argv[1] if len(sys.argv) > 1 else "./all_collected_ds"
+    mode = sys.argv[2] if len(sys.argv) > 2 else "a"
+
+    ds = StreetViewControlDataset(root, mode=mode, image_size=512, augment=True)
+    print(f"Размер датасета: {len(ds)}")
+
+    s = ds[0]
+    print(f"pixel_values:  {s['pixel_values'].shape}  "
+          f"[{s['pixel_values'].min():.2f}, {s['pixel_values'].max():.2f}]")
+    for i, c in enumerate(s["conditioning_pixel_values"]):
+        print(f"control[{i}]:   {c.shape}  [{c.min():.2f}, {c.max():.2f}]")
+    print(f"prompt: {s['prompts']}")
+    print(f"cell:   {s['cell']}")
+

src/depth2street_gen.py

@@ -0,0 +1,216 @@
+import os
+from glob import glob
+from typing import Optional, List, Tuple
+
+import gin
+import torch
+from PIL import Image
+from tqdm import tqdm
+
+from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, EulerAncestralDiscreteScheduler
+
+from utils import ensure_dir, depth_npy_to_pil, list_subdirs
+
+
+@gin.configurable
+class Depth2StreetGenerator:
+    def __init__(
+        self,
+        # --- model settings ---
+        base_model_id: str = "stabilityai/stable-diffusion-xl-base-1.0",
+        controlnet_id: str = "diffusers/controlnet-depth-sdxl-1.0",
+        device: str = "cuda",
+        torch_dtype: str = "fp16",  # "fp16" or "fp32"
+        enable_xformers: bool = False,
+
+        # --- generation settings ---
+        size: int = 1024,
+        steps: int = 30,
+        guidance_scale: float = 5.0,
+        control_scale: float = 1.0,
+        seed: int = 123,
+
+        prompt: str = (
+            "realistic street view photo, eye-level camera, wide angle lens, "
+            "urban street, buildings, road, natural lighting, high detail"
+        ),
+        negative_prompt: str = (
+            "aerial view, top-down, bird's-eye view, map, satellite, isometric, "
+            "low quality, blurry, distortion, warped perspective"
+        ),
+
+        # --- single mode ---
+        depth_npy_path: Optional[str] = None,
+        out_path: str = "out.png",
+
+        # --- dataset mode ---
+        ds_root: Optional[str] = None,
+        depth_filename: str = "depth.npy",
+        out_dir: str = "out_depth2street",
+        out_name: str = "street_gen.png",
+        overwrite: bool = False,
+    ):
+        self.base_model_id = base_model_id
+        self.controlnet_id = controlnet_id
+
+        if device == "cuda" and not torch.cuda.is_available():
+            device = "cpu"
+        self.device = device
+
+        if torch_dtype not in ("fp16", "fp32"):
+            raise ValueError("torch_dtype must be 'fp16' or 'fp32'")
+        self.torch_dtype = torch.float16 if (torch_dtype == "fp16" and self.device.startswith("cuda")) else torch.float32
+
+        self.enable_xformers = bool(enable_xformers)
+
+        self.size = int(size)
+        self.steps = int(steps)
+        self.guidance_scale = float(guidance_scale)
+        self.control_scale = float(control_scale)
+        self.seed = int(seed)
+
+        self.prompt = str(prompt)
+        self.negative_prompt = str(negative_prompt)
+
+        self.depth_npy_path = depth_npy_path
+        self.out_path = out_path
+
+        self.ds_root = ds_root
+        self.depth_filename = depth_filename
+        self.out_dir = out_dir
+        self.out_name = out_name
+        self.overwrite = bool(overwrite)
+
+        self._pipe: Optional[StableDiffusionXLControlNetPipeline] = None
+
+    def _load_pipe(self) -> StableDiffusionXLControlNetPipeline:
+        if self._pipe is not None:
+            return self._pipe
+
+        controlnet = ControlNetModel.from_pretrained(self.controlnet_id, torch_dtype=self.torch_dtype)
+
+        pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+            self.base_model_id,
+            controlnet=controlnet,
+            torch_dtype=self.torch_dtype,
+        )
+
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipe.to(self.device)
+
+        if self.enable_xformers:
+            try:
+                pipe.enable_xformers_memory_efficient_attention()
+            except Exception as e:
+                print(f"[WARN] xFormers enable failed: {e}")
+
+        pipe.set_progress_bar_config(disable=True)
+        self._pipe = pipe
+        return pipe
+
+    def _generate_one(self, depth_npy: str) -> Image.Image:
+        pipe = self._load_pipe()
+        control_img = depth_npy_to_pil(depth_npy, out_size=(self.size, self.size))
+
+        gen = torch.Generator(device=self.device).manual_seed(self.seed)
+
+        img = pipe(
+            prompt=self.prompt,
+            negative_prompt=self.negative_prompt,
+            image=control_img,
+            num_inference_steps=self.steps,
+            guidance_scale=self.guidance_scale,
+            controlnet_conditioning_scale=self.control_scale,
+            generator=gen,
+        ).images[0]
+
+        return img
+
+    def run_single(self) -> None:
+        if not self.depth_npy_path:
+            raise ValueError("depth_npy_path is not set for single mode")
+
+        ensure_dir(os.path.dirname(self.out_path) or ".")
+        img = self._generate_one(self.depth_npy_path)
+        img.save(self.out_path)
+        print(f"[OK] saved: {self.out_path}")
+
+    def run_dataset(self) -> None:
+        if not self.ds_root:
+            raise ValueError("ds_root is not set for dataset mode")
+
+        if not os.path.isdir(self.ds_root):
+            raise FileNotFoundError(f"ds_root not found: {self.ds_root}")
+
+        ensure_dir(self.out_dir)
+
+        # Expected: ds_root/*/depth.npy
+        subdirs = list_subdirs(self.ds_root)
+        if not subdirs:
+            raise RuntimeError(f"No subdirectories found in {self.ds_root}")
+
+        total = 0
+        done = 0
+        skipped = 0
+
+        for sdir in subdirs:
+            depth_path = os.path.join(sdir, self.depth_filename)
+            if not os.path.exists(depth_path):
+                continue
+            total += 1
+
+        if total == 0:
+            raise RuntimeError(f"No '{self.depth_filename}' found under {self.ds_root}/*/")
+
+        print(f"[INFO] Found {total} samples with {self.depth_filename}")
+        print(f"[INFO] Output dir: {self.out_dir} | overwrite={self.overwrite}")
+
+        pbar = tqdm(subdirs, desc="depth2street", unit="sample")
+        for sdir in pbar:
+            depth_path = os.path.join(sdir, self.depth_filename)
+            if not os.path.exists(depth_path):
+                continue
+
+            sample_id = os.path.basename(sdir.rstrip("/"))
+            out_path = os.path.join(self.out_dir, sample_id)
+            ensure_dir(out_path)
+            out_file = os.path.join(out_path, self.out_name)
+
+            if (not self.overwrite) and os.path.exists(out_file):
+                skipped += 1
+                pbar.set_postfix(done=done, skipped=skipped)
+                continue
+
+            img = self._generate_one(depth_path)
+            img.save(out_file)
+            done += 1
+            pbar.set_postfix(done=done, skipped=skipped)
+
+        print(f"[DONE] done={done}, skipped={skipped}, total_with_depth={total}")
+
+    def run(self) -> None:
+        # Автовыбор режима
+        if self.depth_npy_path:
+            self.run_single()
+        elif self.ds_root:
+            self.run_dataset()
+        else:
+            raise ValueError("Set either depth_npy_path (single) or ds_root (dataset)")
+
+
+
+# ! ===========================================================================================
+
+import gin
+
+def main():
+    gin_config_path = os.path.join(os.getcwd(), 'configs/depth2street_single.gin')
+    # gin_config_path = os.path.join(os.getcwd(), 'configs/depth2street_dataset.gin')
+    gin.parse_config_file(gin_config_path)
+
+    gen = Depth2StreetGenerator() 
+    gen.run()
+
+
+if __name__ == "__main__":
+    main()

src/depth2street_onlydrone_depth.py

@@ -0,0 +1,247 @@
+import os
+from typing import Optional, List, Tuple
+
+import gin
+import torch
+from PIL import Image
+from tqdm import tqdm
+
+from diffusers import (
+    ControlNetModel,
+    StableDiffusionXLControlNetPipeline,
+    EulerAncestralDiscreteScheduler,
+)
+from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
+
+from utils import ensure_dir, list_subdirs, load_rgb
+
+
+@gin.configurable
+class DepthEdgeDepth2StreetGenerator:
+    """
+    Experiment C: MultiControlNet with 2 controls:
+      - depth edges image (e.g., depth_drone_edge.png)
+      - depth image (e.g., depth_drone.png)
+
+    Dataset structure expected:
+      ds_root/
+        00000/
+          depth_drone_edge.png
+          depth_drone.png
+          (optional) drone.png, meta.json ...
+        00001/
+          ...
+
+    Output:
+      out_dir/<sample_id>/street_gen.png
+    """
+
+    def __init__(
+        self,
+        # --- models ---
+        base_model_id: str = "stabilityai/stable-diffusion-xl-base-1.0",
+        controlnet_edge_id: str = "diffusers/controlnet-depth-sdxl-1.0",
+        controlnet_depth_id: str = "diffusers/controlnet-depth-sdxl-1.0",
+
+        # --- runtime ---
+        device: str = "cuda",
+        torch_dtype: str = "fp16",  # "fp16" | "fp32"
+        enable_xformers: bool = False,
+
+        # --- data ---
+        ds_root: Optional[str] = None,
+        edge_filename: str = "depth_drone_edge.png",
+        depth_filename: str = "depth_drone.png",
+
+        # --- output ---
+        out_dir: str = "out_street",
+        out_name: str = "street_gen.png",
+        overwrite: bool = False,
+
+        # --- generation ---
+        size: int = 1024,  # SDXL native; можно 768, если VRAM/скорость
+        steps: int = 30,
+        guidance_scale: float = 5.0,
+        seed: int = 123,
+
+        # scales for MultiControlNet: [edge_scale, depth_scale]
+        edge_scale: float = 1.0,
+        depth_scale: float = 0.5,
+
+        prompt: str = (
+            "realistic street-level photo, eye-level viewpoint, camera height 1.6m, "
+            "wide-angle 24mm lens, street, sidewalks, buildings facades, road, "
+            "natural daylight, high detail, sharp focus"
+        ),
+        negative_prompt: str = (
+            "aerial view, drone view, top-down, bird's-eye view, satellite, map, "
+            "isometric, miniature, tilt-shift, orthographic, lowres, blurry, "
+            "warped perspective, distorted geometry"
+        ),
+
+        # optional: process only first N subfolders (0 = all)
+        max_samples: int = 0,
+    ):
+        self.base_model_id = base_model_id
+        self.controlnet_edge_id = controlnet_edge_id
+        self.controlnet_depth_id = controlnet_depth_id
+
+        if device == "cuda" and not torch.cuda.is_available():
+            device = "cpu"
+        self.device = device
+
+        if torch_dtype not in ("fp16", "fp32"):
+            raise ValueError("torch_dtype must be 'fp16' or 'fp32'")
+        self.torch_dtype = torch.float16 if (torch_dtype == "fp16" and self.device.startswith("cuda")) else torch.float32
+        self.enable_xformers = bool(enable_xformers)
+
+        self.ds_root = ds_root
+        self.edge_filename = edge_filename
+        self.depth_filename = depth_filename
+
+        self.out_dir = out_dir
+        self.out_name = out_name
+        self.overwrite = bool(overwrite)
+
+        self.size = int(size)
+        self.steps = int(steps)
+        self.guidance_scale = float(guidance_scale)
+        self.seed = int(seed)
+
+        self.edge_scale = float(edge_scale)
+        self.depth_scale = float(depth_scale)
+
+        self.prompt = str(prompt)
+        self.negative_prompt = str(negative_prompt)
+        self.max_samples = int(max_samples)
+
+        self._pipe: Optional[StableDiffusionXLControlNetPipeline] = None
+
+    def _load_pipe(self) -> StableDiffusionXLControlNetPipeline:
+        if self._pipe is not None:
+            return self._pipe
+
+        cn_edge = ControlNetModel.from_pretrained(self.controlnet_edge_id, torch_dtype=self.torch_dtype)
+        cn_depth = ControlNetModel.from_pretrained(self.controlnet_depth_id, torch_dtype=self.torch_dtype)
+
+        multi_cn = MultiControlNetModel([cn_edge, cn_depth])
+
+        pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+            self.base_model_id,
+            controlnet=multi_cn,
+            torch_dtype=self.torch_dtype,
+        )
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipe.to(self.device)
+
+        # if self.enable_xformers:
+        #     try:
+        #         pipe.enable_xformers_memory_efficient_attention()
+        #     except Exception as e:
+        #         print(f"[WARN] xFormers enable failed: {e}")
+
+        pipe.set_progress_bar_config(disable=True)
+        self._pipe = pipe
+        return pipe
+
+    def _prep_control_img(self, path: str) -> Image.Image:
+        """
+        Load control image and resize to (size,size).
+        ControlNet expects RGB images.
+        """
+        im = load_rgb(path)
+        if im.size != (self.size, self.size):
+            im = im.resize((self.size, self.size), resample=Image.BILINEAR)
+        return im
+
+    def _generate_one(self, edge_path: str, depth_path: str) -> Image.Image:
+        pipe = self._load_pipe()
+
+        edge_img = self._prep_control_img(edge_path)
+        depth_img = self._prep_control_img(depth_path)
+
+        # Important: MultiControlNet expects list of control images aligned with scales
+        control_images = [edge_img, depth_img]
+        control_scales = [self.edge_scale, self.depth_scale]
+
+        gen = torch.Generator(device=self.device).manual_seed(self.seed)
+
+        out = pipe(
+            prompt=self.prompt,
+            negative_prompt=self.negative_prompt,
+            image=control_images,
+            num_inference_steps=self.steps,
+            guidance_scale=self.guidance_scale,
+            controlnet_conditioning_scale=control_scales,
+            generator=gen,
+        )
+        return out.images[0]
+
+    def run(self) -> None:
+        if not self.ds_root:
+            raise ValueError("ds_root is not set")
+        if not os.path.isdir(self.ds_root):
+            raise FileNotFoundError(f"ds_root not found: {self.ds_root}")
+
+        ensure_dir(self.out_dir)
+
+        subdirs = list_subdirs(self.ds_root)
+        if self.max_samples > 0:
+            subdirs = subdirs[: self.max_samples]
+
+        # Pre-scan valid samples
+        samples: List[Tuple[str, str, str]] = []
+        for sdir in subdirs:
+            edge_path = os.path.join(sdir, self.edge_filename)
+            depth_path = os.path.join(sdir, self.depth_filename)
+            if os.path.exists(edge_path) and os.path.exists(depth_path):
+                samples.append((os.path.basename(sdir.rstrip("/")), edge_path, depth_path))
+
+        if not samples:
+            raise RuntimeError(
+                f"No samples found with ({self.edge_filename} AND {self.depth_filename}) under {self.ds_root}/*/"
+            )
+
+        print(f"[INFO] Found {len(samples)} samples")
+        print(f"[INFO] Output dir: {self.out_dir} | overwrite={self.overwrite}")
+        print(f"[INFO] size={self.size} steps={self.steps} guidance={self.guidance_scale} seed={self.seed}")
+        print(f"[INFO] scales: edge={self.edge_scale} depth={self.depth_scale}")
+
+        done = 0
+        skipped = 0
+
+        test = 5
+
+        pbar = tqdm(samples, desc="street-gen (Exp C)", unit="sample")
+        for sample_id, edge_path, depth_path in pbar:
+            out_sdir = os.path.join(self.out_dir, sample_id)
+            ensure_dir(out_sdir)
+            out_file = os.path.join(out_sdir, self.out_name)
+
+            if (not self.overwrite) and os.path.exists(out_file):
+                skipped += 1
+                pbar.set_postfix(done=done, skipped=skipped)
+                continue
+
+            img = self._generate_one(edge_path=edge_path, depth_path=depth_path)
+            img.save(out_file)
+
+            done += 1
+            pbar.set_postfix(done=done, skipped=skipped)
+
+            if done == test:
+                break
+
+        print(f"[DONE] done={done}, skipped={skipped}, total={len(samples)}")
+
+
+def main():
+    gin_config_path = os.path.join(os.getcwd(), 'configs/depth2street_dronly_masks.gin')
+    gin.parse_config_file(gin_config_path)
+
+    gen = DepthEdgeDepth2StreetGenerator()
+    gen.run()
+
+
+if __name__ == "__main__":
+    main()

src/depth_anything_v2_generate.py

@@ -0,0 +1,164 @@
+import os
+from glob import glob
+from typing import List, Optional
+
+import gin
+import numpy as np
+import torch
+from PIL import Image
+import transformers
+from transformers import AutoImageProcessor, AutoModelForDepthEstimation
+
+from utils import ensure_dir, robust_depth_normalize, pil_load_rgb
+
+
+@gin.configurable
+class DepthGenerator:
+    def __init__(
+        self,
+        ds_root: str,
+        source: str = "drone", 
+        model_id: str = "depth-anything/Depth-Anything-V2-Small-hf",
+        device: str = "cuda",
+        out_name: str = "depth.npy",
+        out_vis_name: str = "depth_vis.png",
+        exts: Optional[List[str]] = None,
+        recursive: bool = False,
+        overwrite: bool = False,
+        batch_log_every: int = 50,
+    ):
+        self.ds_root = ds_root
+        self.source = source
+        self.model_id = model_id
+        self.device = device if (device != "cuda" or torch.cuda.is_available()) else "cpu"
+        self.out_name = out_name
+        self.out_vis_name = out_vis_name
+        self.exts = exts if exts is not None else ["png", "jpg", "jpeg"]
+        self.recursive = recursive
+        self.overwrite = overwrite
+        self.batch_log_every = int(batch_log_every)
+
+        if self.source not in ("sat", "drone"):
+            raise ValueError("source must be 'sat' or 'drone'")
+
+    @torch.inference_mode()
+    def predict_depth(
+        self,
+        model: AutoModelForDepthEstimation,
+        processor: AutoImageProcessor,
+        image: Image.Image,
+    ) -> np.ndarray:
+        inputs = processor(images=image, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        outputs = model(**inputs)
+
+        predicted = outputs.predicted_depth  # [B, H', W']
+        predicted = torch.nn.functional.interpolate(
+            predicted.unsqueeze(1),
+            size=(image.height, image.width),
+            mode="bicubic",
+            align_corners=False,
+        ).squeeze(1)  # [B, H, W]
+        depth = predicted[0].detach().float().cpu().numpy().astype(np.float32)
+        return depth
+
+    def save_depth(self, depth: np.ndarray, out_npy: str, out_vis_png: str) -> None:
+        ensure_dir(os.path.dirname(out_npy) if os.path.dirname(out_npy) else ".")
+        np.save(out_npy, depth)
+
+        d01 = robust_depth_normalize(depth)
+        d8 = (d01 * 255.0).astype(np.uint8)
+        Image.fromarray(d8, mode="L").save(out_vis_png)
+
+    def _collect_paths(self) -> List[str]:
+        paths: List[str] = []
+        if self.recursive:
+            for ext in self.exts:
+                patt = os.path.join(self.ds_root, "**", f"{self.source}.{ext}")
+                paths.extend(glob(patt, recursive=True))
+        else:
+            for ext in self.exts:
+                patt = os.path.join(self.ds_root, "*", f"{self.source}.{ext}")
+                paths.extend(glob(patt))
+
+        paths = sorted(list(set(paths)))
+        if not paths:
+            exts_str = ",".join(self.exts)
+            raise RuntimeError(
+                f"No files found. Expected {self.ds_root}/*/{self.source}.({exts_str})"
+            )
+        return paths
+
+    def run(self) -> None:
+        processor, model = self._load_depth_model()
+        model.to(self.device)
+        model.eval()
+
+        paths = self._collect_paths()
+        print(f"[INFO] Found {len(paths)} images for source='{self.source}' under '{self.ds_root}'")
+        print(f"[INFO] Model: {self.model_id} | Device: {self.device}")
+        print(f"[INFO] Output: {self.out_name}, {self.out_vis_name} | overwrite={self.overwrite}")
+
+        done = 0
+        skipped = 0
+
+        for i, p in enumerate(paths, start=1):
+            out_dir = os.path.dirname(p)
+            out_npy = os.path.join(out_dir, self.out_name)
+            out_vis = os.path.join(out_dir, self.out_vis_name)
+
+            if (not self.overwrite) and os.path.exists(out_npy) and os.path.exists(out_vis):
+                skipped += 1
+                continue
+
+            img = pil_load_rgb(p)
+            depth = self.predict_depth(model, processor, img)
+            self.save_depth(depth, out_npy, out_vis)
+            done += 1
+
+            if self.batch_log_every > 0 and (done % self.batch_log_every == 0):
+                print(f"[PROGRESS] done={done} skipped={skipped} / total={len(paths)}")
+
+        print(f"[DONE] done={done}, skipped={skipped}, total={len(paths)}")
+
+    def _load_depth_model(self):
+        try:
+            processor = AutoImageProcessor.from_pretrained(self.model_id)
+            model = AutoModelForDepthEstimation.from_pretrained(self.model_id)
+            return processor, model
+        except Exception as e1:
+            try:
+                processor = AutoImageProcessor.from_pretrained(
+                    self.model_id, trust_remote_code=True
+                )
+                model = AutoModelForDepthEstimation.from_pretrained(
+                    self.model_id, trust_remote_code=True
+                )
+                return processor, model
+            except Exception as e2:
+                msg = str(e1) + " | " + str(e2)
+                if "depth_anything" in msg:
+                    raise RuntimeError(
+                        "Cannot load a Depth-Anything checkpoint with the current "
+                        f"transformers=={transformers.__version__}. "
+                        "Please upgrade Transformers in your active environment, "
+                        'for example: `pip install -U \"transformers>=4.43.0\"`.'
+                    ) from e2
+                raise
+
+
+
+# ! ===========================================================================================
+
+import gin
+
+def main():
+    gin_config_path = os.path.join(os.getcwd(), 'configs/depth_drone.gin')
+    gin.parse_config_file(gin_config_path)
+
+    gen = DepthGenerator() 
+    gen.run()
+
+
+if __name__ == "__main__":
+    main()

src/depth_edges_gen.py

@@ -0,0 +1,135 @@
+import os
+from glob import glob
+from typing import List, Optional
+
+import gin
+import numpy as np
+
+import cv2
+from PIL import Image
+
+
+def ensure_dir(path: str) -> None:
+    if path:
+        os.makedirs(path, exist_ok=True)
+
+
+def robust_normalize(depth: np.ndarray, p_lo: float = 2.0, p_hi: float = 98.0) -> np.ndarray:
+    d = depth.astype(np.float32)
+    lo = np.percentile(d, p_lo)
+    hi = np.percentile(d, p_hi)
+    d = np.clip(d, lo, hi)
+    mn = float(d.min())
+    mx = float(d.max())
+    d01 = (d - mn) / (mx - mn + 1e-6)
+    return np.clip(d01, 0.0, 1.0).astype(np.float32)
+
+
+def compute_edges(depth01: np.ndarray, blur_ksize: int = 5) -> np.ndarray:
+    x = depth01.astype(np.float32)
+
+    if blur_ksize and blur_ksize > 0:
+        # ksize must be odd
+        if blur_ksize % 2 == 0:
+            blur_ksize += 1
+        x = cv2.GaussianBlur(x, (blur_ksize, blur_ksize), 0)
+
+    gx = cv2.Sobel(x, cv2.CV_32F, 1, 0, ksize=3)
+    gy = cv2.Sobel(x, cv2.CV_32F, 0, 1, ksize=3)
+    mag = np.sqrt(gx * gx + gy * gy)
+
+    mag = mag / (float(mag.max()) + 1e-6)
+    return np.clip(mag, 0.0, 1.0).astype(np.float32)
+
+
+def save_png01(x01: np.ndarray, out_path: str) -> None:
+    ensure_dir(os.path.dirname(out_path) or ".")
+    x8 = (np.clip(x01, 0.0, 1.0) * 255.0).astype(np.uint8)
+    Image.fromarray(x8, mode="L").save(out_path)
+
+
+@gin.configurable
+class DepthEdgesGenerator:
+    def __init__(
+        self,
+        ds_root: str,
+        depth_filename: str = "depth.npy",
+        out_name: str = "depth_edge.png",
+
+        # normalization
+        p_lo: float = 2.0,
+        p_hi: float = 98.0,
+
+        # edge settings
+        blur_ksize: int = 5,
+
+        # scan
+        recursive: bool = False,
+        overwrite: bool = False,
+        batch_log_every: int = 50,
+    ):
+        self.ds_root = ds_root
+        self.depth_filename = depth_filename
+        self.out_name = out_name
+
+        self.p_lo = float(p_lo)
+        self.p_hi = float(p_hi)
+        self.blur_ksize = int(blur_ksize)
+
+        self.recursive = bool(recursive)
+        self.overwrite = bool(overwrite)
+        self.batch_log_every = int(batch_log_every)
+
+    def _collect_paths(self) -> List[str]:
+        if self.recursive:
+            patt = os.path.join(self.ds_root, "**", self.depth_filename)
+            paths = glob(patt, recursive=True)
+        else:
+            patt = os.path.join(self.ds_root, "*", self.depth_filename)
+            paths = glob(patt)
+
+        paths = sorted(list(set(paths)))
+        if not paths:
+            raise RuntimeError(f"No files found: {patt}")
+        return paths
+
+    def run(self) -> None:
+        paths = self._collect_paths()
+        print(f"[INFO] Found {len(paths)} depth maps under '{self.ds_root}'")
+        print(f"[INFO] depth_filename={self.depth_filename} out_name={self.out_name} overwrite={self.overwrite}")
+        print(f"[INFO] p_lo={self.p_lo} p_hi={self.p_hi} blur_ksize={self.blur_ksize}")
+
+        done = 0
+        skipped = 0
+
+        for p in paths:
+            sample_dir = os.path.dirname(p)
+            out_path = os.path.join(sample_dir, self.out_name)
+
+            if (not self.overwrite) and os.path.exists(out_path):
+                skipped += 1
+                continue
+
+            depth = np.load(p).astype(np.float32)
+            d01 = robust_normalize(depth, p_lo=self.p_lo, p_hi=self.p_hi)
+            e01 = compute_edges(d01, blur_ksize=self.blur_ksize)
+
+            save_png01(e01, out_path)
+            done += 1
+
+            if self.batch_log_every > 0 and (done % self.batch_log_every == 0):
+                print(f"[PROGRESS] done={done} skipped={skipped} / total={len(paths)}")
+
+        print(f"[DONE] done={done}, skipped={skipped}, total={len(paths)}")
+
+
+def main():
+    gin_config_path = os.path.join(os.getcwd(), 'configs/depth_edge_gen.gin')
+    gin.parse_config_file(gin_config_path)
+
+    gen = DepthEdgesGenerator()
+    gen.run()
+
+
+if __name__ == "__main__":
+    main()

src/train_cn_sdxl_mc_v2.py

@@ -0,0 +1,413 @@
+"""
+train.py — ControlNet SDXL, street view dataset.
+
+Режимы:
+    a  — controls: [drone/depth]
+    b1 — controls: [drone/depth, sat/semantic_color]
+    b2 — controls: [drone/depth, sat/depth]
+    c  — controls: [sat/semantic_color, drone/rgb, street/depth]
+    target всегда: street/forward_rgb
+
+Запуск:
+    accelerate launch train.py \
+        --data_root /mnt/data1tb/CarlaDS/all_collected_ds \
+        --output_dir ./checkpoints \
+        --mode a \
+        --mixed_precision fp16
+"""
+
+import argparse
+import os
+import random
+from typing import List
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from accelerate import Accelerator
+from accelerate.utils import ProjectConfiguration, set_seed
+
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    UNet2DConditionModel,
+    ControlNetModel,
+)
+from transformers import AutoTokenizer, CLIPTextModel, CLIPTextModelWithProjection
+
+from dataset_v2 import StreetViewControlDataset, collate_fn
+from utils import ensure_dir
+
+import bitsandbytes as bnb
+
+# ---------------------------------------------------------------------------
+# Text encoding — SDXL dual encoder
+# ---------------------------------------------------------------------------
+
+def encode_prompt(prompts, tok1, tok2, enc1, enc2, device):
+    def tok(tokenizer, texts):
+        return tokenizer(
+            texts, padding="max_length", truncation=True,
+            max_length=77, return_tensors="pt",
+        ).input_ids.to(device)
+
+    with torch.no_grad():
+        out1 = enc1(tok(tok1, prompts), output_hidden_states=True)
+        out2 = enc2(tok(tok2, prompts), output_hidden_states=True)
+
+    # Penultimate hidden states от обоих энкодеров → конкат [B,77,2048]
+    embeds = torch.cat([out1.hidden_states[-2], out2.hidden_states[-2]], dim=-1)
+    pooled = out2[0]  # [B, 1280]
+    return embeds, pooled
+
+
+# ---------------------------------------------------------------------------
+# Args
+# ---------------------------------------------------------------------------
+
+def parse_args():
+    p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    p.add_argument("--pretrained_model", default="stabilityai/stable-diffusion-xl-base-1.0")
+    p.add_argument("--data_root",   required=True)
+    p.add_argument("--output_dir",  required=True)
+    p.add_argument("--mode",        required=True, choices=["a", "b1", "b2", "c"])
+
+    p.add_argument("--image_size",       type=int,   default=1024)
+    p.add_argument("--sat_seg_blend",    type=float, default=0.5)
+    p.add_argument("--drone_yaw",        default="yaw180")
+    p.add_argument("--street_direction", default="forward")
+    p.add_argument("--prompt_default",
+                   default="realistic street view photo, urban road, high detail")
+    p.add_argument("--prompt_dropout",   type=float, default=0.1)
+    p.add_argument("--augment",          action="store_true")
+
+    p.add_argument("--train_batch_size", type=int,   default=1)
+    p.add_argument("--num_workers",      type=int,   default=4)
+    p.add_argument("--learning_rate",    type=float, default=1e-5)
+    p.add_argument("--max_steps",        type=int,   default=5000)
+    p.add_argument("--grad_accum",       type=int,   default=8)
+    p.add_argument("--mixed_precision",  default="fp16", choices=["no", "fp16", "bf16"])
+    p.add_argument("--save_every",       type=int,   default=500)
+    p.add_argument("--seed",             type=int,   default=42)
+    p.add_argument("--controlnet_init",  default=None,
+                   help="Путь к существующему ControlNet для дообучения")
+
+    return p.parse_args()
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    args = parse_args()
+    ensure_dir(args.output_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.grad_accum,
+        mixed_precision=args.mixed_precision,
+        project_config=ProjectConfiguration(project_dir=args.output_dir),
+    )
+    set_seed(args.seed)
+    device = accelerator.device
+
+    if accelerator.is_main_process:
+        print(f"[Train] mode={args.mode} | device={device} | precision={args.mixed_precision}")
+
+    # ── Dataset ───────────────────────────────────────────────────────────────
+    ds = StreetViewControlDataset(
+        data_root=args.data_root,
+        mode=args.mode,
+        image_size=args.image_size,
+        sat_seg_blend=args.sat_seg_blend,
+        drone_yaw=args.drone_yaw,
+        street_direction=args.street_direction,
+        prompt_default=args.prompt_default,
+        augment=args.augment,
+    )
+    dl = DataLoader(
+        ds,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        drop_last=True,
+    )
+
+    num_controls = ds[0]["num_controls"]
+    if accelerator.is_main_process:
+        print(f"[Dataset] {len(ds)} samples | {num_controls} control(s)")
+
+    # ── Models — все в fp32 ───────────────────────────────────────────────────
+    # HF: "always have all model weights in full float32 precision when starting
+    # training — even if doing mixed precision training"
+    # Accelerator сам управляет autocast и GradScaler через mixed_precision флаг.
+
+    noise_scheduler = DDPMScheduler.from_pretrained(
+        args.pretrained_model, subfolder="scheduler")
+
+    tokenizer_one = AutoTokenizer.from_pretrained(
+        args.pretrained_model, subfolder="tokenizer", use_fast=False)
+    tokenizer_two = AutoTokenizer.from_pretrained(
+        args.pretrained_model, subfolder="tokenizer_2", use_fast=False)
+
+    # variant="fp16" загружает заранее сконвертированные fp16 веса с HF
+    # включая GroupNorm — это обходит проблему mixed dtype
+    text_encoder_one = CLIPTextModel.from_pretrained(
+        args.pretrained_model, subfolder="text_encoder",
+        torch_dtype=torch.float16, variant="fp16")
+    text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
+        args.pretrained_model, subfolder="text_encoder_2",
+        torch_dtype=torch.float16, variant="fp16")
+
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model, subfolder="vae",
+        # torch_dtype=torch.float16, variant="fp16")
+        # torch_dtype=torch.float16, variant="fp32")
+    )
+
+    unet = UNet2DConditionModel.from_pretrained(
+        args.pretrained_model, subfolder="unet",
+        torch_dtype=torch.float16, variant="fp16")
+
+    # Замораживаем всё кроме ControlNet
+    vae.requires_grad_(False)
+    unet.requires_grad_(False)
+    text_encoder_one.requires_grad_(False)
+    text_encoder_two.requires_grad_(False)
+
+    # ── ControlNet ────────────────────────────────────────────────────────────
+    if args.controlnet_init:
+        controlnet = ControlNetModel.from_pretrained(args.controlnet_init)
+    else:
+        controlnet = ControlNetModel.from_unet(unet)
+
+    # Патч первого слоя hint embedding для multi-control
+    if num_controls > 1:
+        old_conv  = controlnet.controlnet_cond_embedding.conv_in
+        new_in_ch = old_conv.in_channels + 3 * (num_controls - 1)
+        new_conv  = torch.nn.Conv2d(
+            new_in_ch, old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+        )
+        with torch.no_grad():
+            new_conv.weight[:, :old_conv.in_channels] = old_conv.weight
+            new_conv.weight[:, old_conv.in_channels:] = 0
+            new_conv.bias.copy_(old_conv.bias)
+        controlnet.controlnet_cond_embedding.conv_in = new_conv
+        if accelerator.is_main_process:
+            print(f"[ControlNet] conv_in patched: {old_conv.in_channels} -> {new_in_ch} ch")
+
+    controlnet.train()
+    controlnet.enable_gradient_checkpointing()
+
+    # ── Optimizer ─────────────────────────────────────────────────────────────
+    # optimizer = torch.optim.AdamW(
+    #     controlnet.parameters(),
+    #     lr=args.learning_rate,
+    #     betas=(0.9, 0.999),
+    #     weight_decay=1e-2,
+    #     eps=1e-8,
+    # )
+
+    optimizer = bnb.optim.AdamW8bit(
+                                controlnet.parameters(), 
+                                lr=args.learning_rate, 
+                                betas=(0.9, 0.999), 
+                                weight_decay=1e-2)
+
+
+    # ── Accelerator prepare ───────────────────────────────────────────────────
+    controlnet, optimizer, dl = accelerator.prepare(controlnet, optimizer, dl)
+
+    vae.to(device, dtype=torch.float32)
+    unet.to(device)
+    text_encoder_one.to(device)
+    text_encoder_two.to(device)
+
+    vae.eval()
+    unet.eval()
+    text_encoder_one.eval()
+    text_encoder_two.eval()
+
+    # ── Training loop ─────────────────────────────────────────────────────────
+    global_step = 0
+    pbar = tqdm(total=args.max_steps,
+                disable=not accelerator.is_main_process,
+                desc="Training")
+
+    while global_step < args.max_steps:
+        for batch in dl:
+            if global_step >= args.max_steps:
+                break
+
+            with accelerator.accumulate(controlnet):
+                pixel_values = batch["pixel_values"].to(device, dtype=torch.float16)
+                controls     = batch["conditioning_pixel_values"].to(device, dtype=torch.float16)
+                prompts      = batch["prompts"]
+
+                # !!!
+                print(f"pixel_values: min={pixel_values.min():.3f} max={pixel_values.max():.3f} nan={pixel_values.isnan().any()} dtype={pixel_values.dtype}")
+
+                # Prompt dropout для CFG
+                if args.prompt_dropout > 0:
+                    prompts = [
+                        "" if random.random() < args.prompt_dropout else p
+                        for p in prompts
+                    ]
+
+                bsz = pixel_values.shape[0]
+
+                # Text embeddings
+                prompt_embeds, pooled_embeds = encode_prompt(
+                    prompts,
+                    tokenizer_one, tokenizer_two,
+                    text_encoder_one, text_encoder_two,
+                    device,
+                )
+
+                # SDXL micro-conditioning time_ids
+                h = w = args.image_size
+                time_ids = torch.tensor(
+                    [h, w, 0, 0, h, w], device=device, dtype=torch.long,
+                ).unsqueeze(0).repeat(bsz, 1)
+
+                added_cond_kwargs = {
+                    "text_embeds": pooled_embeds,
+                    "time_ids":    time_ids,
+                }
+
+                # ! VAE encode
+                # with torch.no_grad():
+                #     latents = vae.encode(pixel_values).latent_dist.sample()
+                #     latents = latents * vae.config.scaling_factor
+
+                # ! check
+                # with torch.autocast("cuda"):
+                #     latents = vae.encode(pixel_values).latent_dist.sample()
+                #     print(f"latents raw: nan={latents.isnan().any()} inf={latents.isinf().any()}")
+                #     latents = latents * vae.config.scaling_factor
+                #     print(f"latents scaled: nan={latents.isnan().any()} scaling_factor={vae.config.scaling_factor}")
+
+                # ! fp32 
+                # with torch.no_grad():
+                #     latents = vae.encode(pixel_values.float()).latent_dist.sample()
+                #     latents = latents * vae.config.scaling_factor
+
+                print("VAE:", next(vae.parameters()).device, next(vae.parameters()).dtype)
+                print("pixel_values:", pixel_values.device, pixel_values.dtype)
+
+                device_type = "cuda" if pixel_values.is_cuda else "cpu"
+
+                with torch.no_grad():
+                    with torch.autocast(device_type=device_type, enabled=False):
+                        pv = pixel_values.float()
+                        latents = vae.encode(pv).latent_dist.sample()
+                        latents = latents * vae.config.scaling_factor
+
+                # weight_dtype для denoiser stack
+                weight_dtype = (
+                    torch.float16 if accelerator.mixed_precision == "fp16"
+                    else torch.bfloat16 if accelerator.mixed_precision == "bf16"
+                    else torch.float32
+                )
+
+                latents = latents.to(dtype=weight_dtype)
+
+                # Diffusion forward process
+                noise     = torch.randn_like(latents)
+                timesteps = torch.randint(
+                    0, noise_scheduler.config.num_train_timesteps,
+                    (bsz,), device=device,
+                ).long()
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # ControlNet + UNet forward в одном autocast блоке.
+                # UNet без torch.no_grad() — градиент должен течь
+                # через down_block_res/mid_block_res обратно в ControlNet.
+                # UNet параметры заморожены (requires_grad=False) —
+                # градиенты через них не накапливаются, только проходят.
+                with torch.autocast("cuda"):
+                    down_block_res, mid_block_res = controlnet(
+                        noisy_latents,
+                        timesteps,
+                        encoder_hidden_states=prompt_embeds,
+                        controlnet_cond=controls,
+                        added_cond_kwargs=added_cond_kwargs,
+                        return_dict=False,
+                    )
+
+                    model_pred = unet(
+                        noisy_latents,
+                        timesteps,
+                        encoder_hidden_states=prompt_embeds,
+                        down_block_additional_residuals=down_block_res,
+                        mid_block_additional_residual=mid_block_res,
+                        added_cond_kwargs=added_cond_kwargs,
+                        return_dict=False,
+                    )[0]
+
+                # Loss в fp32 для стабильности
+                if noise_scheduler.config.prediction_type == "epsilon":
+                    target = noise
+                elif noise_scheduler.config.prediction_type == "v_prediction":
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    raise ValueError(
+                        f"Unknown prediction_type: {noise_scheduler.config.prediction_type}")
+
+                print(f"model_pred: min={model_pred.min():.3f} max={model_pred.max():.3f} nan={model_pred.isnan().any()}")
+                print(f"target:     min={target.min():.3f} max={target.max():.3f} nan={target.isnan().any()}")
+                print(f"controls:   min={controls.min():.3f} max={controls.max():.3f} nan={controls.isnan().any()}")
+                print(f"latents:    min={latents.min():.3f} max={latents.max():.3f} nan={latents.isnan().any()}")
+
+                loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+
+                # if not torch.isfinite(loss):
+                #     if accelerator.is_main_process:
+                #          print(f"[WARN] step {global_step}: loss={loss.item()}, skipping")
+                #     optimizer.zero_grad(set_to_none=True)
+                # else:
+                #     accelerator.backward(loss)
+
+                accelerator.backward(loss)
+
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(controlnet.parameters(), 1.0)
+
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+            if accelerator.is_main_process:
+                pbar.update(1)
+                pbar.set_postfix(step=global_step, loss=f"{loss.item():.4f}")
+
+            global_step += 1
+
+            if accelerator.is_main_process and global_step % args.save_every == 0:
+                _save(accelerator, controlnet, args.output_dir, global_step)
+
+    if accelerator.is_main_process:
+        _save(accelerator, controlnet, args.output_dir, global_step, final=True)
+
+    accelerator.end_training()
+
+
+def _save(accelerator, controlnet, output_dir, step, final=False):
+    tag  = "final" if final else f"checkpoint-{step}"
+    path = os.path.join(output_dir, tag, "controlnet")
+    ensure_dir(path)
+    accelerator.unwrap_model(controlnet).save_pretrained(path)
+    print(f"\n[SAVE] {path}")
+
+
+if __name__ == "__main__":
+    main()
+
+    

src/train_controlnet_sdxl_multicond.py

@@ -0,0 +1,358 @@
+"""
+train.py — ControlNet SDXL, street view dataset.
+
+Режимы:
+    a  — controls: [drone/depth]
+    b1 — controls: [drone/depth, sat/semantic_color]
+    b2 — controls: [drone/depth, sat/depth]
+    c  — controls: [sat/semantic_color, drone/rgb, street/depth]
+    target всегда: street/forward_rgb
+
+Запуск:
+    accelerate launch train.py \
+        --data_root /mnt/data1tb/CarlaDS/all_collected_ds \
+        --output_dir ./checkpoints \
+        --mode a \
+        --mixed_precision fp16
+"""
+
+import argparse
+import os
+import random
+from typing import List
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from accelerate import Accelerator
+from accelerate.utils import ProjectConfiguration, set_seed
+
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    UNet2DConditionModel,
+    ControlNetModel,
+)
+from transformers import AutoTokenizer, CLIPTextModel, CLIPTextModelWithProjection
+
+from dataset_v2 import StreetViewControlDataset, collate_fn
+from utils import ensure_dir
+
+
+# ---------------------------------------------------------------------------
+# Text encoding — SDXL dual encoder
+# ---------------------------------------------------------------------------
+
+def encode_prompt(prompts, tok1, tok2, enc1, enc2, device):
+    def tok(tokenizer, texts):
+        return tokenizer(
+            texts, padding="max_length", truncation=True,
+            max_length=77, return_tensors="pt",
+        ).input_ids.to(device)
+
+    with torch.no_grad():
+        out1 = enc1(tok(tok1, prompts), output_hidden_states=True)
+        out2 = enc2(tok(tok2, prompts), output_hidden_states=True)
+
+    # Penultimate hidden states от обоих энкодеров → конкат [B,77,2048]
+    embeds = torch.cat([out1.hidden_states[-2], out2.hidden_states[-2]], dim=-1)
+    pooled = out2[0]  # [B, 1280]
+    return embeds, pooled
+
+
+# ---------------------------------------------------------------------------
+# Args
+# ---------------------------------------------------------------------------
+
+def parse_args():
+    p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    p.add_argument("--pretrained_model", default="stabilityai/stable-diffusion-xl-base-1.0")
+    p.add_argument("--data_root",   required=True)
+    p.add_argument("--output_dir",  required=True)
+    p.add_argument("--mode",        required=True, choices=["a", "b1", "b2", "c"])
+
+    p.add_argument("--image_size",       type=int,   default=1024)
+    p.add_argument("--sat_seg_blend",    type=float, default=0.5)
+    p.add_argument("--drone_yaw",        default="yaw180")
+    p.add_argument("--street_direction", default="forward")
+    p.add_argument("--prompt_default",
+                   default="realistic street view photo, urban road, high detail")
+    p.add_argument("--prompt_dropout",   type=float, default=0.1)
+    p.add_argument("--augment",          action="store_true")
+
+    p.add_argument("--train_batch_size", type=int,   default=1)
+    p.add_argument("--num_workers",      type=int,   default=4)
+    p.add_argument("--learning_rate",    type=float, default=1e-5)
+    p.add_argument("--max_steps",        type=int,   default=50000)
+    p.add_argument("--grad_accum",       type=int,   default=8)
+    p.add_argument("--mixed_precision",  default="fp16", choices=["no", "fp16", "bf16"])
+    p.add_argument("--save_every",       type=int,   default=2000)
+    p.add_argument("--seed",             type=int,   default=42)
+    p.add_argument("--controlnet_init",  default=None,
+                   help="Путь к существующему ControlNet для дообучения")
+
+    return p.parse_args()
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    args = parse_args()
+    ensure_dir(args.output_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.grad_accum,
+        mixed_precision=args.mixed_precision,
+        project_config=ProjectConfiguration(project_dir=args.output_dir),
+    )
+    set_seed(args.seed)
+    device = accelerator.device
+
+    if accelerator.is_main_process:
+        print(f"[Train] mode={args.mode} | device={device} | precision={args.mixed_precision}")
+
+    # ── Dataset ───────────────────────────────────────────────────────────────
+    ds = StreetViewControlDataset(
+        data_root=args.data_root,
+        mode=args.mode,
+        image_size=args.image_size,
+        sat_seg_blend=args.sat_seg_blend,
+        drone_yaw=args.drone_yaw,
+        street_direction=args.street_direction,
+        prompt_default=args.prompt_default,
+        augment=args.augment,
+    )
+    dl = DataLoader(
+        ds,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        drop_last=True,
+    )
+
+    num_controls = ds[0]["num_controls"]
+    if accelerator.is_main_process:
+        print(f"[Dataset] {len(ds)} samples | {num_controls} control(s)")
+
+    # ── Models — все в fp32 ───────────────────────────────────────────────────
+    # HF: "always have all model weights in full float32 precision when starting
+    # training — even if doing mixed precision training"
+    # Accelerator сам управляет autocast и GradScaler через mixed_precision флаг.
+
+    noise_scheduler = DDPMScheduler.from_pretrained(
+        args.pretrained_model, subfolder="scheduler")
+
+    tokenizer_one = AutoTokenizer.from_pretrained(
+        args.pretrained_model, subfolder="tokenizer", use_fast=False)
+    tokenizer_two = AutoTokenizer.from_pretrained(
+        args.pretrained_model, subfolder="tokenizer_2", use_fast=False)
+
+    text_encoder_one = CLIPTextModel.from_pretrained(
+        args.pretrained_model, subfolder="text_encoder")
+    text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
+        args.pretrained_model, subfolder="text_encoder_2")
+
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model, subfolder="vae")
+
+    unet = UNet2DConditionModel.from_pretrained(
+        args.pretrained_model, subfolder="unet")
+
+    # Замораживаем всё кроме ControlNet
+    vae.requires_grad_(False)
+    unet.requires_grad_(False)
+    text_encoder_one.requires_grad_(False)
+    text_encoder_two.requires_grad_(False)
+
+    # ── ControlNet ────────────────────────────────────────────────────────────
+    if args.controlnet_init:
+        controlnet = ControlNetModel.from_pretrained(args.controlnet_init)
+    else:
+        controlnet = ControlNetModel.from_unet(unet)
+
+    # Патч первого слоя hint embedding для multi-control
+    if num_controls > 1:
+        old_conv  = controlnet.controlnet_cond_embedding.conv_in
+        new_in_ch = old_conv.in_channels + 3 * (num_controls - 1)
+        new_conv  = torch.nn.Conv2d(
+            new_in_ch, old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+        )
+        with torch.no_grad():
+            new_conv.weight[:, :old_conv.in_channels] = old_conv.weight
+            new_conv.weight[:, old_conv.in_channels:] = 0
+            new_conv.bias.copy_(old_conv.bias)
+        controlnet.controlnet_cond_embedding.conv_in = new_conv
+        if accelerator.is_main_process:
+            print(f"[ControlNet] conv_in patched: {old_conv.in_channels} -> {new_in_ch} ch")
+
+    controlnet.train()
+    controlnet.enable_gradient_checkpointing()
+
+    # ── Optimizer ─────────────────────────────────────────────────────────────
+    optimizer = torch.optim.AdamW(
+        controlnet.parameters(),
+        lr=args.learning_rate,
+        betas=(0.9, 0.999),
+        weight_decay=1e-2,
+        eps=1e-8,
+    )
+
+    # ── Accelerator prepare ───────────────────────────────────────────────────
+    # Передаём только обучаемые объекты — замороженные модели переносим вручную.
+    # Если передать замороженные модели в prepare, accelerator может обернуть их
+    # так что граф вычислений не строится и backward падает.
+    controlnet, optimizer, dl = accelerator.prepare(controlnet, optimizer, dl)
+
+    # Замороженные модели вручную на device в weight_dtype
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    vae.to(device, dtype=weight_dtype)
+    unet.to(device)  # UNet в fp32 — GroupNorm несовместим с fp16
+    text_encoder_one.to(device, dtype=weight_dtype)
+    text_encoder_two.to(device, dtype=weight_dtype)
+
+    vae.eval()
+    unet.eval()
+    text_encoder_one.eval()
+    text_encoder_two.eval()
+
+    # ── Training loop ─────────────────────────────────────────────────────────
+    global_step = 0
+    pbar = tqdm(total=args.max_steps,
+                disable=not accelerator.is_main_process,
+                desc="Training")
+
+    while global_step < args.max_steps:
+        for batch in dl:
+            if global_step >= args.max_steps:
+                break
+
+            with accelerator.accumulate(controlnet):
+                pixel_values = batch["pixel_values"].to(device, dtype=weight_dtype)
+                controls     = batch["conditioning_pixel_values"].to(device, dtype=weight_dtype)
+                prompts      = batch["prompts"]
+
+                # Prompt dropout для CFG
+                if args.prompt_dropout > 0:
+                    prompts = [
+                        "" if random.random() < args.prompt_dropout else p
+                        for p in prompts
+                    ]
+
+                bsz = pixel_values.shape[0]
+
+                # Text embeddings
+                prompt_embeds, pooled_embeds = encode_prompt(
+                    prompts,
+                    tokenizer_one, tokenizer_two,
+                    text_encoder_one, text_encoder_two,
+                    device,
+                )
+
+                # SDXL micro-conditioning time_ids
+                h = w = args.image_size
+                time_ids = torch.tensor(
+                    [h, w, 0, 0, h, w], device=device, dtype=torch.long,
+                ).unsqueeze(0).repeat(bsz, 1)
+
+                added_cond_kwargs = {
+                    "text_embeds": pooled_embeds,
+                    "time_ids":    time_ids,
+                }
+
+                # VAE encode
+                with torch.no_grad():
+                    latents = vae.encode(pixel_values).latent_dist.sample()
+                    latents = latents * vae.config.scaling_factor
+
+                # Diffusion forward process
+                noise     = torch.randn_like(latents)
+                timesteps = torch.randint(
+                    0, noise_scheduler.config.num_train_timesteps,
+                    (bsz,), device=device,
+                ).long()
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # ControlNet forward (градиенты сохраняются)
+                down_block_res, mid_block_res = controlnet(
+                    noisy_latents,
+                    timesteps,
+                    encoder_hidden_states=prompt_embeds,
+                    controlnet_cond=controls,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )
+
+                # UNet forward в fp32 (без градиентов)
+                with torch.no_grad():
+                    model_pred = unet(
+                        noisy_latents.float(),
+                        timesteps,
+                        encoder_hidden_states=prompt_embeds.float(),
+                        down_block_additional_residuals=[r.float() for r in down_block_res],
+                        mid_block_additional_residual=mid_block_res.float(),
+                        added_cond_kwargs={
+                            "text_embeds": pooled_embeds.float(),
+                            "time_ids":    time_ids,
+                        },
+                        return_dict=False,
+                    )[0]
+
+                # Loss
+                if noise_scheduler.config.prediction_type == "epsilon":
+                    target = noise
+                elif noise_scheduler.config.prediction_type == "v_prediction":
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    raise ValueError(
+                        f"Unknown prediction_type: {noise_scheduler.config.prediction_type}")
+
+                loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+
+                accelerator.backward(loss)
+
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(controlnet.parameters(), 1.0)
+
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+            if accelerator.is_main_process:
+                pbar.update(1)
+                pbar.set_postfix(step=global_step, loss=f"{loss.item():.4f}")
+
+            global_step += 1
+
+            if accelerator.is_main_process and global_step % args.save_every == 0:
+                _save(accelerator, controlnet, args.output_dir, global_step)
+
+    if accelerator.is_main_process:
+        _save(accelerator, controlnet, args.output_dir, global_step, final=True)
+
+    accelerator.end_training()
+
+
+def _save(accelerator, controlnet, output_dir, step, final=False):
+    tag  = "final" if final else f"checkpoint-{step}"
+    path = os.path.join(output_dir, tag, "controlnet")
+    ensure_dir(path)
+    accelerator.unwrap_model(controlnet).save_pretrained(path)
+    print(f"\n[SAVE] {path}")
+
+
+if __name__ == "__main__":
+    main()

src/unet_mem_test.py

@@ -0,0 +1,14 @@
+import torch
+from diffusers import ControlNetModel, UNet2DConditionModel
+
+unet = UNet2DConditionModel.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    subfolder="unet", torch_dtype=torch.float16
+)
+cn = ControlNetModel.from_unet(unet)
+unet.to("cpu")
+
+params = sum(p.numel() for p in cn.parameters())
+print(f"ControlNet params: {params/1e6:.0f}M")
+print(f"VRAM estimate fp16: {params*2/1e9:.1f} GB")
+

src/utils.py

@@ -0,0 +1,91 @@
+import os
+import random
+from dataclasses import dataclass
+from typing import Optional, List, Tuple, Dict, Any
+import numpy as np
+import torch
+from PIL import Image
+
+
+@dataclass
+class SamplePaths:
+    street: str
+    sat: str
+    seg: str
+    drone: Optional[str] = None
+    depth_npy: Optional[str] = None
+    meta_json: Optional[str] = None
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def ensure_dir(path: str) -> None:
+    os.makedirs(path, exist_ok=True)
+
+
+def list_subdirs(root: str) -> List[str]:
+    out = []
+    for name in sorted(os.listdir(root)):
+        p = os.path.join(root, name)
+        if os.path.isdir(p):
+            out.append(p)
+    return out 
+
+
+def safe_join(base: str, rel: str) -> str:
+    return os.path.normpath(os.path.join(base, rel))
+
+
+def pil_load_rgb(path: str) -> Image.Image:
+    return Image.open(path).convert("RGB")
+
+
+def pil_load_l(path: str) -> Image.Image:
+    return Image.open(path).convert("L")
+
+
+def save_image_grid(images: List[Image.Image], rows: int, cols: int, out_path: str, pad: int = 8) -> None:
+    assert len(images) == rows * cols
+    w, h = images[0].size
+    grid_w = cols * w + (cols - 1) * pad
+    grid_h = rows * h + (rows - 1) * pad
+    grid = Image.new("RGB", (grid_w, grid_h), (0, 0, 0))
+    for idx, im in enumerate(images):
+        r = idx // cols
+        c = idx % cols
+        x = c * (w + pad)
+        y = r * (h + pad)
+        grid.paste(im, (x, y))
+    ensure_dir(os.path.dirname(out_path) if os.path.dirname(out_path) else ".")
+    grid.save(out_path)
+
+
+def robust_depth_normalize(depth: np.ndarray) -> np.ndarray:
+    """
+    depth: float32 HxW (любая шкала) -> float32 HxW в [0,1] (робастно по процентилям)
+    """
+    if depth.ndim != 2:
+        raise ValueError(f"Expected depth 2D array HxW, got shape={depth.shape}")
+    d = depth.astype(np.float32)
+    lo = np.percentile(d, 2.0)
+    hi = np.percentile(d, 98.0)
+    d = (d - lo) / (hi - lo + 1e-6)
+    d = np.clip(d, 0.0, 1.0)
+    return d
+
+    
+def depth_npy_to_pil(depth_path: str, out_size: Optional[Tuple[int, int]] = None) -> Image.Image:
+    """Load depth.npy (HxW float32) -> PIL RGB control image."""
+    depth = np.load(depth_path).astype(np.float32)
+    d01 = robust_depth_normalize(depth)
+    d8 = (d01 * 255.0).astype(np.uint8)
+    img = Image.fromarray(d8, mode="L").convert("RGB")
+    if out_size is not None:
+        img = img.resize(out_size, resample=Image.BILINEAR)
+    return img
+

src/validate_samples.py

@@ -0,0 +1,110 @@
+import argparse
+import os
+from typing import List
+
+import torch
+from PIL import Image
+
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, MultiControlNetModel
+from transformers import AutoTokenizer
+
+from dataset import StreetViewControlDataset
+from utils import ensure_dir, save_image_grid
+
+
+def tensor_to_pil(x: torch.Tensor) -> Image.Image:
+    """
+    x: Tensor(3,H,W) in [0,1]
+    """
+    x = x.detach().cpu().clamp(0, 1)
+    x = (x * 255.0).byte()
+    x = x.permute(1, 2, 0).numpy()
+    return Image.fromarray(x, mode="RGB")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--pretrained_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0")
+    parser.add_argument("--controlnet_dir", type=str, required=True,
+                        help="Папка с сохранённым controlnet (например out/.../final/controlnet)")
+    parser.add_argument("--data_root", type=str, required=True)
+    parser.add_argument("--mode", type=str, choices=["a", "b", "c"], required=True)
+    parser.add_argument("--out_dir", type=str, default="val_out")
+    parser.add_argument("--num_samples", type=int, default=4)
+    parser.add_argument("--steps", type=int, default=30)
+    parser.add_argument("--guidance_scale", type=float, default=5.0)
+    parser.add_argument("--control_scales", type=str, default="1.0",
+                        help='Напр: "1.0" или "1.0,0.8,1.2" (по числу контролов)')
+    parser.add_argument("--seed", type=int, default=123)
+    parser.add_argument("--image_size", type=int, default=1024)
+    parser.add_argument("--sat_seg_blend", type=float, default=0.5)
+    parser.add_argument("--prompt_default", type=str, default="realistic street view photo")
+    args = parser.parse_args()
+
+    ensure_dir(args.out_dir)
+
+    ds = StreetViewControlDataset(
+        data_root=args.data_root,
+        mode=args.mode,
+        image_size=args.image_size,
+        sat_seg_blend=args.sat_seg_blend,
+        prompt_default=args.prompt_default,
+    )
+
+    # Determine number of controls from dataset
+    num_controls = len(ds[0]["conditioning_pixel_values"])
+    scales = [float(s.strip()) for s in args.control_scales.split(",")]
+    if len(scales) == 1 and num_controls > 1:
+        scales = scales * num_controls
+    if len(scales) != num_controls:
+        raise ValueError(f"control_scales length must be 1 or == num_controls ({num_controls}), got {len(scales)}")
+
+    # Load ControlNet
+    # В train мы сохраняли MultiControlNetModel как один объект (save_pretrained). Это работает и для одиночного.
+    controlnet = ControlNetModel.from_pretrained(args.controlnet_dir)
+    # Если у вас num_controls>1, то save_pretrained() у MultiControlNetModel сохраняет папку
+    # со списком net-ов; корректная загрузка в diffusers — MultiControlNetModel.from_pretrained().
+    if num_controls > 1:
+        controlnet = MultiControlNetModel.from_pretrained(args.controlnet_dir)
+
+    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+        args.pretrained_model,
+        controlnet=controlnet,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe.to(device)
+    pipe.set_progress_bar_config(disable=True)
+
+    g = torch.Generator(device=device).manual_seed(args.seed)
+
+    images: List[Image.Image] = []
+    n = min(args.num_samples, len(ds))
+    for i in range(n):
+        item = ds[i]
+        prompt = item["prompt"]
+        controls = item["conditioning_pixel_values"]  # list[Tensor(3,H,W)] in [0,1]
+
+        ctrl_pils = [tensor_to_pil(c) for c in controls]
+        result = pipe(
+            prompt=prompt,
+            image=ctrl_pils if num_controls > 1 else ctrl_pils[0],
+            num_inference_steps=args.steps,
+            guidance_scale=args.guidance_scale,
+            controlnet_conditioning_scale=scales if num_controls > 1 else scales[0],
+            generator=g,
+        ).images[0]
+
+        # Сохраним также контрольные входы рядом (для дебага)
+        images.append(result)
+
+    # grid: 1 row x n cols
+    grid_path = os.path.join(args.out_dir, f"grid_{args.mode}.png")
+    save_image_grid(images, rows=1, cols=len(images), out_path=grid_path)
+    print(f"[OK] Saved {grid_path}")
+
+
+if __name__ == "__main__":
+    main()
+
+