Update shape_e_service.py

shape_e_service.py

@@ -1,6 +1,7 @@
 import gc
 import os
-import tempfile
+import uuid
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
 import gradio as gr
@@ -13,6 +14,11 @@ pipe = None
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 
+ROOT_DIR = Path(__file__).resolve().parent
+DATA_DIR = ROOT_DIR / "data"
+ASSETS_DIR = DATA_DIR / "assets"
+ASSETS_DIR.mkdir(parents=True, exist_ok=True)
+
 EXAMPLES = [
     "A cute stylized robot with a round head",
     "A fantasy treasure chest with gold trim",
@@ -37,13 +43,17 @@ def get_pipeline():
     return pipe
 
 
-def save_frames_to_files(frames) -> List[str]:
+def save_frames_to_files(frames, prompt: str) -> List[str]:
+    asset_id = f"asset_{uuid.uuid4().hex[:8]}"
+    asset_dir = ASSETS_DIR / asset_id
+    asset_dir.mkdir(parents=True, exist_ok=True)
+
     frame_paths = []
-    for frame in frames:
-        fd, frame_path = tempfile.mkstemp(suffix=".png", prefix="shape_frame_")
-        os.close(fd)
+    for i, frame in enumerate(frames):
+        frame_path = asset_dir / f"view_{i:03d}.png"
         frame.save(frame_path)
-        frame_paths.append(frame_path)
+        frame_paths.append(str(frame_path))
+
     return frame_paths
 
 
@@ -136,7 +146,7 @@ def generate_and_add_asset(
     )
 
     frames = result.images[0]
-    frame_paths = save_frames_to_files(frames)
+    frame_paths = save_frames_to_files(frames, prompt)
 
     new_asset = make_asset(prompt, frame_paths)
     saved_assets = saved_assets + [new_asset]