Update video_service.py

video_service.py

@@ -71,6 +71,8 @@ def log_tensor_info(tensor, name="Tensor"):
         print("  - O tensor está vazio, sem estatísticas.")
     print("------------------------------------------\n")
 
+
+
 # --- 5. CLASSE PRINCIPAL DO SERVIÇO ---
 class VideoService:
     def __init__(self):
@@ -78,6 +80,9 @@ class VideoService:
         self.config = self._load_config()
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.last_memory_reserved_mb = 0
+        self._tmp_dirs = set()
+        self._tmp_files = set()
+        self._last_outputs = []
         self.pipeline, self.latent_upsampler = self._load_models()
         print(f"Movendo modelos para o dispositivo de inferência: {self.device}")
         self.pipeline.to(self.device)
@@ -88,6 +93,63 @@ class VideoService:
             self._log_gpu_memory("Após carregar modelos")
         print("VideoService pronto para uso.")
     
+    def _register_tmp_dir(self, d: str):
+        if d and os.path.isdir(d):
+            self._tmp_dirs.add(d)
+
+    def _register_tmp_file(self, f: str):
+        if f and os.path.isfile(f):
+            self._tmp_files.add(f)
+
+    def finalize(self, keep_paths=None, extra_paths=None, clear_gpu=True):
+        """
+        Remove temporários e coleta memória. 
+        keep_paths: caminhos que não devem ser removidos (ex.: vídeo final).
+        extra_paths: caminhos adicionais para tentar remover (opcional).
+        """
+        keep = set(keep_paths or [])
+        extras = set(extra_paths or [])
+
+        # Remoção de arquivos
+        for f in list(self._tmp_files | extras):
+            try:
+                if f not in keep and os.path.isfile(f):
+                    os.remove(f)
+            except Exception:
+                pass
+            finally:
+                self._tmp_files.discard(f)
+
+        # Remoção de diretórios
+        for d in list(self._tmp_dirs):
+            try:
+                if d not in keep and os.path.isdir(d):
+                    shutil.rmtree(d, ignore_errors=True)
+            except Exception:
+                pass
+            finally:
+                self._tmp_dirs.discard(d)
+
+        # Coleta de GC e limpeza de VRAM
+        gc.collect()
+        try:
+            import torch
+            if clear_gpu and torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                # Limpa buffers de IPC quando aplicável
+                try:
+                    torch.cuda.ipc_collect()
+                except Exception:
+                    pass
+        except Exception:
+            pass
+
+        # Log opcional pós-limpeza
+        try:
+            self._log_gpu_memory("Após finalize")
+        except Exception:
+            pass
+    
     def _query_gpu_processes_via_nvml(device_index: int) -> List[Dict]:
         try:
             import psutil
@@ -328,6 +390,8 @@ class VideoService:
                     progress_callback(i + 1, total_frames)
         
         self._log_gpu_memory("Fim da Geração")
+
+        finalize()
         return output_video_path, used_seed
 
 print("Criando instância do VideoService. O carregamento do modelo começará agora...")