Update video_service.py

video_service.py

@@ -1,4 +1,3 @@
-
 # video_service.py
 
 import torch
@@ -30,7 +29,7 @@ def run_setup():
         print(f"ERRO CRÍTICO DURANTE O SETUP: 'setup.py' falhou com código {e.returncode}.")
         sys.exit(1)
 
-DEPS_DIR = Path("/data")
+DEPS_DIR = Path("./deps")
 LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
 if not LTX_VIDEO_REPO_DIR.exists():
     run_setup()
@@ -41,7 +40,7 @@ def add_deps_to_path():
     if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
         sys.path.insert(0, str(LTX_VIDEO_REPO_DIR.resolve()))
 
-#add_deps_to_path()
+add_deps_to_path()
 
 # Importações específicas do modelo
 from inference import (
@@ -123,163 +122,141 @@ class VideoService:
                     return worker
             time.sleep(0.1)
 
-    
-    def generate(self, prompt, negative_prompt, input_image_filepath=None, input_video_filepath=None, height=512, width=704, mode="text-to-video", duration=2.0, frames_to_use=9, seed=42, randomize_seed=True, guidance_scale=1.0,  # Agora usado corretamente
+    def generate(self, prompt, negative_prompt, input_image_filepath=None, input_video_filepath=None,
+                 height=512, width=704, mode="text-to-video", duration=2.0,
+                 frames_to_use=9, seed=42, randomize_seed=True, guidance_scale=1.0, # Ignorado, mas mantido por compatibilidade
                  improve_texture=True, progress_callback=None):
         
+        # A MÁGICA DO LAZY LOADING ACONTECE AQUI
         self._ensure_models_are_loaded()
+        
         worker = self._acquire_worker()
         base_device = worker['devices']['base']
         upscaler_device = worker['devices']['upscaler']
         
         try:
-            # Validações alinhadas com app-20.py
-            if mode == "image-to-video" and not input_image_filepath:
-                raise ValueError("Caminho da imagem obrigatório para o modo image-to-video")
-            if mode == "video-to-video" and not input_video_filepath:
-                raise ValueError("Caminho do vídeo obrigatório para o modo video-to-video")
-            
+            # ... (todo o resto do código da função generate permanece exatamente o mesmo) ...
+            if mode == "image-to-video" and not input_image_filepath: raise ValueError("Caminho da imagem é obrigatório para o modo image-to-video")
+            if mode == "video-to-video" and not input_video_filepath: raise ValueError("Caminho do vídeo é obrigatório para o modo video-to-video")
+
             used_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
             seed_everething(used_seed)
-            
-            FPS = 30.0  # Alinhado com app-20.py
-            MAX_NUM_FRAMES = 257
-            target_frames_ideal = duration * FPS
-            target_frames_rounded = round(target_frames_ideal)
-            if target_frames_rounded < 1: target_frames_rounded = 1
-            n_val = round(float(target_frames_rounded - 1.0) / 8.0)
+
+            FPS = 24.0; MAX_NUM_FRAMES = 257
+            target_frames_rounded = round(duration * FPS)
+            n_val = round((float(target_frames_rounded) - 1.0) / 8.0)
             actual_num_frames = max(9, min(MAX_NUM_FRAMES, int(n_val * 8 + 1)))
             
-            actual_height = int(height)
-            actual_width = int(width)
-            height_padded = (actual_height - 1) // 32 * 32 + 32
-            width_padded = (actual_width - 1) // 32 * 32 + 32
-            num_frames_padded = (actual_num_frames - 2) // 8 * 8 + 1  # Alinhamento exato com app-20.py
-            if num_frames_padded != actual_num_frames:
-                print(f"Warning: actual_num_frames {actual_num_frames} and num_frames_padded {num_frames_padded} differ. Using num_frames_padded for pipeline.")
-            
-            padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded)
+            height_padded = ((height - 1) // 32 + 1) * 32
+            width_padded = ((width - 1) // 32 + 1) * 32
+            padding_values = calculate_padding(height, width, height_padded, width_padded)
             pad_left, pad_right, pad_top, pad_bottom = padding_values
             
-            # Kwargs base alinhados
-            call_kwargs = {
-                "prompt": prompt,
-                "negative_prompt": negative_prompt,
-                "height": height_padded,
-                "width": width_padded,
-                "num_frames": num_frames_padded,
-                "framerate": int(FPS),
-                "generator": torch.Generator(device=base_device).manual_seed(used_seed),
-                "output_type": "pt",
-                "conditioning_items": None,
-                "media_items": None,
-                "decode_timestep": self.config['decode_timestep'],
-                "decode_noise_scale": self.config['decode_noise_scale'],
-                "stochastic_sampling": self.config['stochastic_sampling'],
-                "image_cond_noise_scale": 0.15,  # Alinhado
-                "is_video": True,
-                "vae_per_channel_normalize": True,
-                "mixed_precision": self.config['precision'] + " mixed_precision",
-                "offload_to_cpu": False,
-                "enhance_prompt": False,
+            call_kwargs_base = {
+                "prompt": prompt, "negative_prompt": negative_prompt, "num_frames": actual_num_frames, "frame_rate": int(FPS),
+                "decode_timestep": 0.05, "decode_noise_scale": self.config["decode_noise_scale"],
+                "stochastic_sampling": self.config["stochastic_sampling"], "image_cond_noise_scale": 0.025,
+                "is_video": True, "vae_per_channel_normalize": True, "mixed_precision": (self.config["precision"] == "mixed_precision"),
+                "offload_to_cpu": False, "enhance_prompt": False, "skip_layer_strategy": SkipLayerStrategy.AttentionValues
             }
-            
-            # Estratégia de skip layer alinhada
-            stg_mode_str = self.config.get('stg_mode', 'attention_values')
-            if stg_mode_str.lower() in ['stgav', 'attentionvalues']:
-                call_kwargs['skip_layer_strategy'] = SkipLayerStrategy.AttentionValues
-            # ... (adicionar outros elif como no app-20.py)
-            
-            # Conditioning para modos
-            if mode == "image-to-video" and input_image_filepath:
-                media_tensor = load_image_to_tensor_with_resize_and_crop(input_image_filepath, actual_height, actual_width)
-                media_tensor = torch.nn.functional.pad(media_tensor, padding_values)
-                call_kwargs['conditioning_items'] = ConditioningItem(media_tensor.to(base_device), 0, 1.0)
-            elif mode == "video-to-video" and input_video_filepath:
-                call_kwargs['media_items'] = load_media_file(media_path=input_video_filepath, height=actual_height, width=actual_width, max_frames=int(frames_to_use), padding=padding_values).to(base_device)
-            
-            result_images_tensor = None
+
+            result_tensor = None
             if improve_texture:
-                # Alinhamento exato: Use LTXMultiScalePipeline como no app-20.py
-                active_latent_upsampler = worker['latent_upsampler']
-                if not active_latent_upsampler:
-                    raise ValueError("Spatial upscaler model not loaded or improve_texture not selected, cannot use multi-scale.")
-                
-                multi_scale_pipeline_obj = LTXMultiScalePipeline(worker['base_pipeline'], active_latent_upsampler)
-                
-                first_pass_args = self.config.get('first_pass', {}).copy()
-                first_pass_args['guidance_scale'] = float(guidance_scale)  # Override UI
-                first_pass_args.pop('num_inference_steps', None)
+                downscale_factor = self.config.get("downscale_factor", 0.5)
+                downscaled_height_ideal = int(height_padded * downscale_factor); downscaled_width_ideal = int(width_padded * downscale_factor)
+                downscaled_height = ((downscaled_height_ideal - 1) // 32 + 1) * 32; downscaled_width = ((downscaled_width_ideal - 1) // 32 + 1) * 32
+
+                # --- PASSE 1 ---
+                first_pass_kwargs = call_kwargs_base.copy()
+                first_pass_kwargs.update({
+                    "height": downscaled_height, "width": downscaled_width,
+                    "generator": torch.Generator(device=base_device).manual_seed(used_seed),
+                    "output_type": "latent", "guidance_scale": 1.0,
+                    "timesteps": self.config["first_pass"]["timesteps"],
+                    "stg_scale": self.config["first_pass"]["stg_scale"],
+                    "rescaling_scale": self.config["first_pass"]["rescaling_scale"],
+                    "skip_block_list": self.config["first_pass"]["skip_block_list"]
+                })
+
+                if mode == "image-to-video":
+                    padding_low_res = calculate_padding(downscaled_height, downscaled_width, downscaled_height, downscaled_width)
+                    media_tensor_low_res = load_image_to_tensor_with_resize_and_crop(input_image_filepath, downscaled_height, downscaled_width)
+                    media_tensor_low_res = torch.nn.functional.pad(media_tensor_low_res, padding_low_res)
+                    first_pass_kwargs["conditioning_items"] = [ConditioningItem(media_tensor_low_res.to(base_device), 0, 1.0)]
+
+                print(f"Worker {worker['id']}: Iniciando passe 1 em {base_device}")
+                with torch.no_grad(): low_res_latents = worker['base_pipeline'](**first_pass_kwargs).images
                 
-                second_pass_args = self.config.get('second_pass', {}).copy()
-                second_pass_args['guidance_scale'] = float(guidance_scale)  # Override UI
-                second_pass_args.pop('num_inference_steps', None)
+                low_res_latents = low_res_latents.to(upscaler_device)
+                with torch.no_grad(): high_res_latents = worker['latent_upsampler'](low_res_latents)
+                high_res_latents = high_res_latents.to(base_device)
                 
-                multi_scale_call_kwargs = call_kwargs.copy()
-                multi_scale_call_kwargs.update({
-                    "downscale_factor": self.config['downscale_factor'],
-                    "first_pass": first_pass_args,
-                    "second_pass": second_pass_args,
+                # --- PASSE 2 ---
+                second_pass_kwargs = call_kwargs_base.copy()
+                high_res_h, high_res_w = downscaled_height * 2, downscaled_width * 2
+                second_pass_kwargs.update({
+                    "height": high_res_h, "width": high_res_w, "latents": high_res_latents,
+                    "generator": torch.Generator(device=base_device).manual_seed(used_seed),
+                    "output_type": "pt", "image_cond_noise_scale": 0.0, "guidance_scale": 1.0,
+                    "timesteps": self.config["second_pass"]["timesteps"],
+                    "stg_scale": self.config["second_pass"]["stg_scale"],
+                    "rescaling_scale": self.config["second_pass"]["rescaling_scale"],
+                    "skip_block_list": self.config["second_pass"]["skip_block_list"],
+                    "tone_map_compression_ratio": self.config["second_pass"].get("tone_map_compression_ratio", 0.0)
                 })
                 
-                print(f"Calling multi-scale pipeline eff. HxW {actual_height}x{actual_width}, Frames {actual_num_frames} - Padded {num_frames_padded} on {base_device}")
-                result_images_tensor = multi_scale_pipeline_obj(**multi_scale_call_kwargs).images
-            else:
-                # Single-pass alinhado
-                single_pass_call_kwargs = call_kwargs.copy()
-                first_pass_config_from_yaml = self.config.get('first_pass', {})
-                single_pass_call_kwargs['timesteps'] = first_pass_config_from_yaml.get('timesteps')
-                single_pass_call_kwargs['guidance_scale'] = float(guidance_scale)  # Override UI
-                single_pass_call_kwargs['stg_scale'] = first_pass_config_from_yaml.get('stg_scale')
-                single_pass_call_kwargs['rescaling_scale'] = first_pass_config_from_yaml.get('rescaling_scale')
-                single_pass_call_kwargs['skip_block_list'] = first_pass_config_from_yaml.get('skip_block_list')
-                single_pass_call_kwargs.pop('num_inference_steps', None)
-                single_pass_call_kwargs.pop('first_pass', None)
-                single_pass_call_kwargs.pop('second_pass', None)
-                single_pass_call_kwargs.pop('downscale_factor', None)
+                if mode == "image-to-video":
+                    padding_high_res = calculate_padding(high_res_h, high_res_w, high_res_h, high_res_w)
+                    media_tensor_high_res = load_image_to_tensor_with_resize_and_crop(input_image_filepath, high_res_h, high_res_w)
+                    media_tensor_high_res = torch.nn.functional.pad(media_tensor_high_res, padding_high_res)
+                    second_pass_kwargs["conditioning_items"] = [ConditioningItem(media_tensor_high_res.to(base_device), 0, 1.0)]
+
+                print(f"Worker {worker['id']}: Iniciando passe 2 em {base_device}")
+                with torch.no_grad(): result_tensor = worker['base_pipeline'](**second_pass_kwargs).images
+
+            else: # Passe Único
+                single_pass_kwargs = call_kwargs_base.copy()
+                first_pass_config = self.config["first_pass"]
+                single_pass_kwargs.update({
+                    "height": height_padded, "width": width_padded, "output_type": "pt",
+                    "generator": torch.Generator(device=base_device).manual_seed(used_seed),
+                    "guidance_scale": 1.0, **first_pass_config
+                })
+                if mode == "image-to-video":
+                    media_tensor_final = load_image_to_tensor_with_resize_and_crop(input_image_filepath, height_padded, width_padded)
+                    media_tensor_final = torch.nn.functional.pad(media_tensor_final, padding_values)
+                    single_pass_kwargs["conditioning_items"] = [ConditioningItem(media_tensor_final.to(base_device), 0, 1.0)]
+                elif mode == "video-to-video":
+                    single_pass_kwargs["media_items"] = load_media_file(media_path=input_video_filepath, height=height_padded, width=width_padded, max_frames=int(frames_to_use), padding=padding_values).to(base_device)
                 
-                print(f"Calling base pipeline padded HxW {height_padded}x{width_padded}, Frames {actual_num_frames} - Padded {num_frames_padded} on {base_device}")
-                result_images_tensor = worker['base_pipeline'](**single_pass_call_kwargs).images
-            
-            if result_images_tensor is None:
-                raise ValueError("Generation failed.")
-            
-            # Slicing e salvamento alinhados
-            slice_h_end = -pad_bottom if pad_bottom > 0 else None
-            slice_w_end = -pad_right if pad_right > 0 else None
-            result_images_tensor = result_images_tensor[:, :, :actual_num_frames, pad_top:slice_h_end, pad_left:slice_w_end]
-            video_np = result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy()
-            video_np = np.clip(video_np, 0, 1) * 255.0
-            video_np = video_np.astype(np.uint8)
+                print(f"Worker {worker['id']}: Iniciando passe único em {base_device}")
+                with torch.no_grad(): result_tensor = worker['base_pipeline'](**single_pass_kwargs).images
             
+            if result_tensor.shape[-2:] != (height, width):
+                num_frames_final = result_tensor.shape[2]
+                videos_tensor = result_tensor.permute(0, 2, 1, 3, 4).reshape(-1, result_tensor.shape[1], result_tensor.shape[3], result_tensor.shape[4])
+                videos_resized = torch.nn.functional.interpolate(videos_tensor, size=(height, width), mode='bilinear', align_corners=False)
+                result_tensor = videos_resized.reshape(result_tensor.shape[0], num_frames_final, result_tensor.shape[1], height, width).permute(0, 2, 1, 3, 4)
+
+            result_tensor = result_tensor[:, :, :actual_num_frames, (pad_top if pad_top > 0 else None):(-pad_bottom if pad_bottom > 0 else None), (pad_left if pad_left > 0 else None):(-pad_right if pad_right > 0 else None)]
+            video_np = (result_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).astype(np.uint8)
             temp_dir = tempfile.mkdtemp()
             output_video_path = os.path.join(temp_dir, f"output_{used_seed}.mp4")
-            try:
-                with imageio.get_writer(output_video_path, fps=call_kwargs['framerate'], macro_block_size=1) as video_writer:
-                    for frame_idx in range(video_np.shape[0]):
-                        if progress_callback:
-                            progress_callback(frame_idx / video_np.shape[0], desc="Saving video")
-                        video_writer.append_data(video_np[frame_idx])
-            except Exception as e:
-                print(f"Error saving video with macro_block_size=1: {e}")
-                with imageio.get_writer(output_video_path, fps=call_kwargs['framerate'], format='FFMPEG', codec='libx264', quality=8) as video_writer:
-                    for frame_idx in range(video_np.shape[0]):
-                        if progress_callback:
-                            progress_callback(frame_idx / video_np.shape[0], desc="Saving video fallback ffmpeg")
-                        video_writer.append_data(video_np[frame_idx])
-            
+
+            with imageio.get_writer(output_video_path, fps=call_kwargs_base["frame_rate"], codec='libx264', quality=8) as writer:
+                for i, frame in enumerate(video_np):
+                    writer.append_data(frame)
+                    if progress_callback: progress_callback(i + 1, len(video_np))
             return output_video_path, used_seed
-        
+
         except Exception as e:
-            print(f"!!!!!!!! ERRO no Worker {worker['id']}: {e} !!!!!!!!")
+            print(f"!!!!!!!! ERRO no Worker {worker['id']} !!!!!!!!\n{e}")
             raise e
         finally:
-            print(f"Worker {worker['id']} Tarefa finalizada. Limpando cache e liberando worker...")
-            with torch.cuda.device(base_device):
-                torch.cuda.empty_cache()
-            with torch.cuda.device(upscaler_device):
-                torch.cuda.empty_cache()
-            worker['lock'].release()
+            print(f"Worker {worker['id']}: Tarefa finalizada. Limpando cache e liberando worker...")
+            with torch.cuda.device(base_device): torch.cuda.empty_cache()
+            with torch.cuda.device(upscaler_device): torch.cuda.empty_cache()
+            worker["lock"].release()
 
 # A instância do serviço é criada aqui, mas os modelos só serão carregados no primeiro clique.
-video_generation_service = VideoService()
+video_generation_service = VideoService()