reorder

tour-comfy/edit_api.py

@@ -1366,14 +1366,10 @@ def _write_all_static() -> None:
             _write_json(os.path.join(data_dir, "group-names.json"),
                         database.get_all_group_names())
 
-            # videos.json — wireframe template videos
+            # videos.json — wireframe template videos and grouped assets
             wireframe_dir = _load_wireframe_dir()
-            videos = []
-            if os.path.isdir(wireframe_dir):
-                videos = [f for f in sorted(os.listdir(wireframe_dir))
-                          if f.lower().endswith(VIDEO_EXTENSIONS) and not f.startswith('.')]
             _write_json(os.path.join(data_dir, "videos.json"),
-                        {"videos": videos, "wireframe_dir": wireframe_dir})
+                        _get_grouped_wireframes(wireframe_dir))
 
             # config.json — current config snapshot
             try:
@@ -2001,17 +1997,81 @@ def list_images(archived: bool = False):
         raise HTTPException(500, str(e))
 
 
+def _get_grouped_wireframes(wireframe_dir: str) -> dict:
+    """Scan and return a grouped structure of all videos, clips, and frames in wireframe_dir."""
+    if not os.path.isdir(wireframe_dir):
+        return {"videos": [], "groups": [], "standalone_images": []}
+        
+    all_files = sorted(os.listdir(wireframe_dir))
+    videos = [f for f in all_files if f.lower().endswith(VIDEO_EXTENSIONS) and not f.startswith('.')]
+    images = [f for f in all_files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')) and not f.startswith('.')]
+    
+    stem_map = {}
+    
+    # Match primary videos and trimmed clips
+    for v in videos:
+        stem = os.path.splitext(v)[0]
+        # Match e.g. dance_10s-20s.mp4
+        m = re.match(r"^(.*)_(\d+)s-(\d+)s$", stem)
+        if m:
+            parent_stem = m.group(1)
+            stem_map.setdefault(parent_stem, {"video": None, "clips": [], "frames": []})
+            stem_map[parent_stem]["clips"].append(v)
+        else:
+            stem_map.setdefault(stem, {"video": None, "clips": [], "frames": []})
+            stem_map[stem]["video"] = v
+            
+    # Match extracted frames
+    standalone_images = []
+    for img in images:
+        stem = os.path.splitext(img)[0]
+        # Match e.g. dance-f120.png
+        m = re.match(r"^(.*)-f(\d+)$", stem)
+        if m:
+            parent_stem = m.group(1)
+            stem_map.setdefault(parent_stem, {"video": None, "clips": [], "frames": []})
+            stem_map[parent_stem]["frames"].append(img)
+        else:
+            standalone_images.append(img)
+            
+    # Format and sort
+    groups = []
+    for stem, data in sorted(stem_map.items()):
+        # Sort clips by start time if possible
+        def get_clip_start(c):
+            stem_c = os.path.splitext(c)[0]
+            m_c = re.match(r"^.*_(\d+)s-(\d+)s$", stem_c)
+            return int(m_c.group(1)) if m_c else 0
+            
+        # Sort frames numerically by frame number
+        def get_frame_num(f):
+            stem_f = os.path.splitext(f)[0]
+            m_f = re.match(r"^.*-f(\d+)$", stem_f)
+            return int(m_f.group(1)) if m_f else 0
+            
+        data["clips"].sort(key=get_clip_start)
+        data["frames"].sort(key=get_frame_num)
+        
+        groups.append({
+            "stem": stem,
+            "video": data["video"],
+            "clips": data["clips"],
+            "frames": data["frames"]
+        })
+        
+    return {
+        "videos": videos,
+        "groups": groups,
+        "standalone_images": standalone_images,
+        "wireframe_dir": wireframe_dir
+    }
+
+
 @app.get("/videos")
 def list_videos():
-    """Return available wireframe template videos."""
+    """Return available wireframe template videos and grouped wireframe assets."""
     wireframe_dir = _load_wireframe_dir()
-    if not os.path.isdir(wireframe_dir):
-        return {"videos": []}
-    videos = [
-        f for f in sorted(os.listdir(wireframe_dir))
-        if f.lower().endswith(VIDEO_EXTENSIONS) and not f.startswith('.')
-    ]
-    return {"videos": videos, "wireframe_dir": wireframe_dir}
+    return _get_grouped_wireframes(wireframe_dir)
 
 
 @app.get("/wireframe/frame/{video_name}")
@@ -2216,9 +2276,31 @@ def wireframe_extract_frame(req: FrameExtractRequest):
         img = _extract_frame_at(video_path, req.time)
     except Exception as e:
         raise HTTPException(500, str(e))
+    
+    # Calculate saved filename to return to client
+    try:
+        import cv2
+        cap = cv2.VideoCapture(video_path)
+        if cap.isOpened():
+            fps = cap.get(cv2.CAP_PROP_FPS)
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            cap.release()
+            frame_num = int(round(req.time * fps)) if fps > 0 else 0
+            frame_num = max(0, min(total - 1, frame_num))
+        else:
+            frame_num = int(round(req.time * 30.0))
+    except Exception:
+        frame_num = int(round(req.time * 30.0))
+        
+    stem = os.path.splitext(req.video_name)[0]
+    saved_filename = f"{stem}-f{frame_num}.png"
+    
     buf = io.BytesIO()
     img.save(buf, format="PNG")
-    return {"frame_b64": base64.b64encode(buf.getvalue()).decode()}
+    return {
+        "frame_b64": base64.b64encode(buf.getvalue()).decode(),
+        "filename": saved_filename
+    }
 
 
 class FaceswapRequest(BaseModel):
@@ -2734,6 +2816,29 @@ def upload_image(
     return {"status": "processing", "filename": filename, "group_id": effective_gid, "prompts": prompt_list}
 
 
+@app.post("/wireframe/upload")
+def upload_wireframe_image(
+    image: UploadFile = File(...),
+):
+    wireframe_dir = _load_wireframe_dir()
+    os.makedirs(wireframe_dir, exist_ok=True)
+
+    ts = time.strftime("%Y%m%d_%H%M%S")
+    safe_filename = re.sub(r'[^a-zA-Z0-9_.-]', '_', image.filename or "uploaded")
+    # Ensure extension
+    if not safe_filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
+        safe_filename += ".png"
+
+    filename = f"up_{ts}_{safe_filename}"
+    file_path = os.path.join(wireframe_dir, filename)
+
+    with open(file_path, "wb") as f:
+        shutil.copyfileobj(image.file, f)
+
+    _invalidate_static()
+    return {"status": "success", "filename": filename}
+
+
 @app.post("/edit")
 async def edit(
     image: UploadFile = File(...),
@@ -3381,7 +3486,7 @@ def remove_background_group(group_id: str, background_tasks: BackgroundTasks):
 # --- scenery generation -------------------------------------------------------
 
 def _extract_frame_at(video_path: str, t: float) -> Image.Image:
-    """Extract a single frame at time t (seconds) from a video via ffmpeg."""
+    """Extract a single frame at time t (seconds) from a video via ffmpeg, and save to wireframe dir."""
     import subprocess as _sp
     r = _sp.run(
         ['ffmpeg', '-y', '-ss', str(t), '-i', video_path,
@@ -3390,7 +3495,36 @@ def _extract_frame_at(video_path: str, t: float) -> Image.Image:
     )
     if r.returncode != 0 or not r.stdout:
         raise ValueError(f"ffmpeg frame extract failed: {r.stderr.decode(errors='replace')[:300]}")
-    return Image.open(io.BytesIO(r.stdout)).convert("RGB")
+    img = Image.open(io.BytesIO(r.stdout)).convert("RGB")
+
+    # Save to wireframe directory as videoname-f{frame_number}.png
+    try:
+        import cv2
+        cap = cv2.VideoCapture(video_path)
+        if cap.isOpened():
+            fps = cap.get(cv2.CAP_PROP_FPS)
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            cap.release()
+            frame_num = int(round(t * fps)) if fps > 0 else 0
+            frame_num = max(0, min(total - 1, frame_num))
+        else:
+            frame_num = int(round(t * 30.0))
+    except Exception as e:
+        print(f"Error computing frame number: {e}")
+        frame_num = int(round(t * 30.0))
+
+    try:
+        wireframe_dir = os.path.dirname(video_path)
+        video_name = os.path.basename(video_path)
+        stem = os.path.splitext(video_name)[0]
+        out_name = f"{stem}-f{frame_num}.png"
+        out_path = os.path.join(wireframe_dir, out_name)
+        img.save(out_path, format="PNG")
+        print(f"[scenery] Saved extracted frame to {out_path}")
+    except Exception as e:
+        print(f"Error saving extracted frame: {e}")
+
+    return img
 
 
 class SceneryRequest(BaseModel):
@@ -3398,6 +3532,7 @@ class SceneryRequest(BaseModel):
     scene_bytes: str | None = None  # base64-encoded PNG/JPEG of the reference scene → image1
     scene_video: str | None = None  # wireframe video name to extract frame from
     scene_time: float = 0.0       # timestamp (seconds) to extract from video
+    scene_image: str | None = None  # optional wireframe image name to use directly
     extra_filename: str | None = None  # optional extra reference in output_dir → image3 (Picture 3)
     prompt: str | None = None     # override; auto-built if None
     seed: int = -1
@@ -3420,7 +3555,8 @@ def _make_side_by_side(img1: Image.Image, img2: Image.Image,
 
 def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
                     prompt: str, seed: int, extra_pils: list | None = None,
-                    scene_video: str | None = None, extra_filename: str | None = None):
+                    scene_video: str | None = None, scene_image: str | None = None,
+                    extra_filename: str | None = None):
     output_dir = _load_output_dir()
     try:
         model_path = os.path.join(output_dir, model_filename)
@@ -3457,6 +3593,8 @@ def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
             refs = [model_filename]
             if scene_video:
                 refs.append(f"video:{scene_video}")
+            if scene_image:
+                refs.append(f"wireframe:{scene_image}")
             if extra_filename:
                 refs.append(extra_filename)
             database.upsert_person(out_name, filepath=out_path, embedding=embedding,
@@ -3491,6 +3629,11 @@ def generate_scenery(req: SceneryRequest):
         import base64
         raw = base64.b64decode(req.scene_bytes)
         scene_pil = Image.open(io.BytesIO(raw)).convert("RGB")
+    elif req.scene_image:
+        image_path = os.path.join(wireframe_dir, req.scene_image)
+        if not os.path.exists(image_path):
+            raise HTTPException(404, f"Scene image not found: {req.scene_image}")
+        scene_pil = Image.open(image_path).convert("RGB")
     elif req.scene_video:
         video_path = os.path.join(wireframe_dir, req.scene_video)
         if not os.path.exists(video_path):
@@ -3500,7 +3643,7 @@ def generate_scenery(req: SceneryRequest):
         except Exception as e:
             raise HTTPException(500, f"Frame extraction failed: {e}")
     else:
-        raise HTTPException(400, "Provide scene_bytes or scene_video")
+        raise HTTPException(400, "Provide scene_bytes, scene_image, or scene_video")
 
     # Optional image3 (Picture 3) — extra reference from output_dir
     extra_pils = []
@@ -3525,6 +3668,7 @@ def generate_scenery(req: SceneryRequest):
         target=_scenery_worker,
         args=(job_id, req.model_filename, scene_pil, prompt, req.seed, extra_pils),
         kwargs={"scene_video": req.scene_video if req.scene_video else None,
+                "scene_image": req.scene_image if req.scene_image else None,
                 "extra_filename": req.extra_filename},
         daemon=True,
     ).start()
@@ -3558,6 +3702,15 @@ def scenery_library():
         except Exception:
             refs = []
         video = next((r[len("video:"):] for r in refs if r.startswith("video:")), None)
+        if not video:
+            wf_img = next((r[len("wireframe:"):] for r in refs if r.startswith("wireframe:")), None)
+            if wf_img:
+                m = re.match(r"^(.*)-f\d+\.(png|jpg|jpeg|webp)$", wf_img, re.IGNORECASE)
+                if m:
+                    video = m.group(1)
+                else:
+                    video = wf_img
+        
         entry = {"filename": filename, "group_id": group_id, "refs": refs}
         if video:
             by_video.setdefault(video, []).append(entry)