diff --git a/rename_children.py b/rename_children.py
index a710768..47dd79c 100755
--- a/rename_children.py
+++ b/rename_children.py
@@ -5,36 +5,48 @@ import dataclasses
 import datetime
 import json
 import logging
-from pathlib import Path
 import re
+from pathlib import Path
 
 
 def parse_args():
-    arger = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    arger.add_argument('dirs', help='Path to the directories', type=Path, nargs='+')
+    arger = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    arger.add_argument("dirs", help="Path to the directories", type=Path, nargs="+")
+    arger.add_argument(
+        "--flatten", help="Flatten directory structure", action="store_true"
+    )
     return arger.parse_args()
 
 
-def natural_sort_key(path: Path) -> tuple[str, ...]:
-    s = path.stem
-    return tuple(f'{int(p):010d}' if p.isnumeric() else p for p in re.findall(r'(\D+|\d+)', s))
+def rename_children(source_dir: Path, flatten: bool) -> None:
+    def natural_sort_key(path: Path) -> tuple[str, ...]:
+        s = path.stem
+        return tuple(
+            f"{int(p):010d}" if p.isnumeric() else p
+            for p in re.findall(r"(\D+|\d+)", s)
+        )
 
+    assert source_dir.is_dir()
 
-def rename_children(cwd: Path) -> None:
-    assert cwd.is_dir()
+    if flatten:
+        files = [f for f in source_dir.rglob("*") if f.is_file()]
+    else:
+        files = [f for f in source_dir.iterdir() if f.is_file()]
 
-    files = [f for f in cwd.glob('*') if f.is_file()]
     files = sorted(files, key=natural_sort_key)
 
-    new_stem = cwd.name
-
     jobs = []
     for i, f in enumerate(files):
-        target = f.with_stem(f'{new_stem}__{i:04d}')
-        if f.suffix == '.jpeg':
-            target = target.with_suffix('.jpg')
+        ext = ".jpg" if f.suffix.lower() in {".jpeg", ".jpg"} else f.suffix
+        target_dir = f.parent
+        if flatten:
+            target_dir = source_dir
+        new_stem = source_dir.name
+        target = target_dir / f"{new_stem}__{i:04d}{ext}"
         if target.is_file():
-            logging.error(f'target already exists: {target}')
+            logging.error(f"target already exists: {target}")
             continue
         jobs.append(MoveOp(source=f, target=target))
         f.rename(target)
@@ -52,24 +64,24 @@ class MoveOp:
 
     def to_dict(self) -> dict:
         return {
-            'source': str(self.source),
-            'target': str(self.target),
+            "source": str(self.source),
+            "target": str(self.target),
         }
 
 
 def write_undo_script(save_dir: Path, ops: list[MoveOp]):
-    now = datetime.datetime.now().isoformat().replace(':', '')
-    save_path = save_dir / f'undo_{now}.json'
-    with save_path.open('w', encoding='utf-8') as f:
+    now = datetime.datetime.now().isoformat().replace(":", "")
+    save_path = save_dir / f"undo_{now}.json"
+    with save_path.open("w", encoding="utf-8") as f:
         for it in ops:
-            f.write(json.dumps(it.to_dict()) + '\n')
+            f.write(json.dumps(it.to_dict()) + "\n")
 
 
 def main():
     args = parse_args()
     for d in args.dirs:
-        rename_children(d)
+        rename_children(source_dir=d, flatten=args.flatten)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()