EuuIia commited on
Commit
8ce4529
·
verified ·
1 Parent(s): ac23084

Update video_service.py

Browse files
Files changed (1) hide show
  1. video_service.py +189 -222
video_service.py CHANGED
@@ -1,6 +1,5 @@
1
  # video_service.py
2
 
3
- # --- 1. IMPORTAÇÕES ---
4
  import torch
5
  import numpy as np
6
  import random
@@ -9,15 +8,15 @@ import yaml
9
  from pathlib import Path
10
  import imageio
11
  import tempfile
12
- from huggingface_hub import hf_hub_download
13
  import sys
14
  import subprocess
15
- from PIL import Image
 
 
16
 
17
- # --- 2. GERENCIAMENTO DE DEPENDÊNCIAS E SETUP ---
18
 
19
  def run_setup():
20
- """Executa o script setup.py para clonar as dependências necessárias."""
21
  setup_script_path = "setup.py"
22
  if not os.path.exists(setup_script_path):
23
  print("AVISO: script 'setup.py' não encontrado. Pulando a clonagem de dependências.")
@@ -36,7 +35,6 @@ if not LTX_VIDEO_REPO_DIR.exists():
36
  run_setup()
37
 
38
  def add_deps_to_path():
39
- """Adiciona o repositório clonado ao sys.path para que suas bibliotecas possam ser importadas."""
40
  if not LTX_VIDEO_REPO_DIR.exists():
41
  raise FileNotFoundError(f"Repositório LTX-Video não encontrado em '{LTX_VIDEO_REPO_DIR}'. Execute o setup.")
42
  if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
@@ -44,252 +42,221 @@ def add_deps_to_path():
44
 
45
  add_deps_to_path()
46
 
47
- # --- 3. IMPORTAÇÕES ESPECÍFICAS DO MODELO ---
48
  from inference import (
49
  create_ltx_video_pipeline, create_latent_upsampler,
50
  load_image_to_tensor_with_resize_and_crop, seed_everething,
51
  calculate_padding, load_media_file
52
  )
53
- from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline
54
  from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
55
 
56
- # --- 4. FUNÇÕES HELPER DE LOG ---
57
- def log_tensor_info(tensor, name="Tensor"):
58
- if not isinstance(tensor, torch.Tensor):
59
- print(f"\n[INFO] O item '{name}' não é um tensor para logar.")
60
- return
61
- print(f"\n--- Informações do Tensor: {name} ---")
62
- print(f" - Shape: {tensor.shape}")
63
- print(f" - Dtype: {tensor.dtype}")
64
- print(f" - Device: {tensor.device}")
65
- if tensor.numel() > 0:
66
- print(f" - Min valor: {tensor.min().item():.4f}")
67
- print(f" - Max valor: {tensor.max().item():.4f}")
68
- print(f" - Média: {tensor.mean().item():.4f}")
69
- else:
70
- print(" - O tensor está vazio, sem estatísticas.")
71
- print("------------------------------------------\n")
72
 
73
- # --- 5. CLASSE PRINCIPAL DO SERVIÇO ---
74
  class VideoService:
75
  def __init__(self):
76
- print("Inicializando VideoService...")
 
 
77
  self.config = self._load_config()
78
- self.device = "cuda" if torch.cuda.is_available() else "cpu"
79
- self.last_memory_reserved_mb = 0
80
- self.pipeline, self.latent_upsampler = self._load_models()
81
- print(f"Movendo modelos para o dispositivo de inferência: {self.device}")
82
- self.pipeline.to(self.device)
83
- if self.latent_upsampler:
84
- self.latent_upsampler.to(self.device)
85
- if self.device == "cuda":
86
- torch.cuda.empty_cache()
87
- self._log_gpu_memory("Após carregar modelos")
88
- print("VideoService pronto para uso.")
89
 
90
- def _log_gpu_memory(self, stage_name: str):
91
- if self.device != "cuda": return
92
- current_reserved_b = torch.cuda.memory_reserved()
93
- current_reserved_mb = current_reserved_b / (1024 ** 2)
94
- total_memory_b = torch.cuda.get_device_properties(0).total_memory
95
- total_memory_mb = total_memory_b / (1024 ** 2)
96
- peak_reserved_mb = torch.cuda.max_memory_reserved() / (1024 ** 2)
97
- delta_mb = current_reserved_mb - self.last_memory_reserved_mb
98
- print(f"\n--- [LOG DE MEMÓRIA GPU] - {stage_name} ---")
99
- print(f" - Uso Atual (Reservado): {current_reserved_mb:.2f} MB / {total_memory_mb:.2f} MB")
100
- print(f" - Variação desde o último log: {delta_mb:+.2f} MB")
101
- if peak_reserved_mb > self.last_memory_reserved_mb:
102
- print(f" - Pico de Uso (nesta operação): {peak_reserved_mb:.2f} MB")
103
- print("--------------------------------------------------\n")
104
- self.last_memory_reserved_mb = current_reserved_mb
105
 
106
  def _load_config(self):
107
  config_file_path = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled.yaml"
108
  with open(config_file_path, "r") as file:
109
  return yaml.safe_load(file)
110
 
111
- def _load_models(self):
112
- models_dir = "downloaded_models_gradio"
113
- Path(models_dir).mkdir(parents=True, exist_ok=True)
114
  LTX_REPO = "Lightricks/LTX-Video"
115
- distilled_model_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["checkpoint_path"], local_dir=models_dir, local_dir_use_symlinks=False)
116
- self.config["checkpoint_path"] = distilled_model_path
117
- spatial_upscaler_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["spatial_upscaler_model_path"], local_dir=models_dir, local_dir_use_symlinks=False)
118
- self.config["spatial_upscaler_model_path"] = spatial_upscaler_path
119
- pipeline = create_ltx_video_pipeline(ckpt_path=self.config["checkpoint_path"], precision=self.config["precision"], text_encoder_model_name_or_path=self.config["text_encoder_model_name_or_path"], sampler=self.config["sampler"], device="cpu", enhance_prompt=False, prompt_enhancer_image_caption_model_name_or_path=self.config["prompt_enhancer_image_caption_model_name_or_path"], prompt_enhancer_llm_model_name_or_path=self.config["prompt_enhancer_llm_model_name_or_path"])
120
- latent_upsampler = None
121
- if self.config.get("spatial_upscaler_model_path"):
122
- latent_upsampler = create_latent_upsampler(self.config["spatial_upscaler_model_path"], device="cpu")
123
- return pipeline, latent_upsampler
124
-
125
- def _prepare_conditioning_tensor_from_file(self, filepath, height, width, padding_values):
126
- """Prepara um tensor de condicionamento a partir de um arquivo de imagem."""
127
- tensor = load_image_to_tensor_with_resize_and_crop(filepath, height, width)
128
- tensor = torch.nn.functional.pad(tensor, padding_values)
129
- return tensor.to(self.device)
130
 
131
- def _extract_frames_from_video(self, video_path: str, frame_indices: list) -> list:
132
- print(f"[INFO] Extraindo frames nos índices: {frame_indices} do vídeo '{video_path}'")
133
- extracted_frames = []
134
- indices_to_get = set(frame_indices)
135
- try:
136
- with imageio.get_reader(video_path) as reader:
137
- for i, frame in enumerate(reader):
138
- if i in indices_to_get:
139
- extracted_frames.append(frame)
140
- if len(extracted_frames) == len(indices_to_get):
141
- break
142
- if len(extracted_frames) != len(frame_indices):
143
- print(f"[AVISO] Esperava extrair {len(frame_indices)} frames, mas o vídeo só tinha {len(extracted_frames)} correspondentes.")
144
- except Exception as e:
145
- print(f"[ERRO] Falha ao extrair frames do vídeo: {e}")
146
- return extracted_frames
147
-
148
- def _get_video_dimensions(self, video_path: str) -> tuple[int, int]:
149
- """Lê um arquivo de vídeo e retorna sua largura e altura."""
150
- try:
151
- with imageio.get_reader(video_path) as reader:
152
- meta = reader.get_meta_data()
153
- size = meta.get('size')
154
- if size:
155
- return size
156
- return (None, None)
157
- except Exception as e:
158
- print(f"[ERRO] Não foi possível ler as dimensões do vídeo: {e}")
159
- return (None, None)
160
-
161
- def generate(self, prompt, negative_prompt, mode="text-to-video",
162
- start_image_filepath=None,
163
- middle_image_filepath=None, middle_frame_number=None, middle_image_weight=1.0,
164
- end_image_filepath=None, end_image_weight=1.0,
165
- input_video_filepath=None, height=512, width=704, duration=2.0,
166
- frames_to_use=9, seed=42, randomize_seed=True, guidance_scale=3.0,
167
- improve_texture=True, progress_callback=None):
168
- if self.device == "cuda":
169
- torch.cuda.empty_cache()
170
- torch.cuda.reset_peak_memory_stats()
171
- self._log_gpu_memory("Início da Geração")
172
 
173
- if mode == "image-to-video" and not start_image_filepath:
174
- raise ValueError("A imagem de início é obrigatória para o modo image-to-video")
175
- if mode == "video-to-video" and not input_video_filepath:
176
- raise ValueError("O vídeo de entrada é obrigatório para o modo video-to-video")
 
 
 
177
 
178
- used_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
179
- seed_everething(used_seed)
 
 
 
 
 
180
 
181
- if mode == "video-to-video":
182
- orig_w, orig_h = self._get_video_dimensions(input_video_filepath)
183
- if orig_w and orig_h:
184
- width = round(orig_w / 32) * 32
185
- height = round(orig_h / 32) * 32
186
- print(f"[INFO] Modo video-to-video: Dimensões recalculadas para {width}x{height}")
187
-
188
- FPS = 24.0
189
- MAX_NUM_FRAMES = 257
190
- target_frames_rounded = round(duration * FPS)
191
- n_val = round((float(target_frames_rounded) - 1.0) / 8.0)
192
- actual_num_frames = max(9, min(MAX_NUM_FRAMES, int(n_val * 8 + 1)))
193
-
194
- height_padded = ((height - 1) // 32 + 1) * 32
195
- width_padded = ((width - 1) // 32 + 1) * 32
196
- padding_values = calculate_padding(height, width, height_padded, width_padded)
197
-
198
- generator = torch.Generator(device=self.device).manual_seed(used_seed)
199
 
200
- conditioning_items = []
 
201
 
202
- if mode == "image-to-video":
203
- start_tensor = self._prepare_conditioning_tensor_from_file(start_image_filepath, height, width, padding_values)
204
- conditioning_items.append(ConditioningItem(start_tensor, 0, 1.0))
205
- if middle_image_filepath and middle_frame_number is not None:
206
- middle_tensor = self._prepare_conditioning_tensor_from_file(middle_image_filepath, height, width, padding_values)
207
- safe_middle_frame = max(0, min(int(middle_frame_number), actual_num_frames - 1))
208
- conditioning_items.append(ConditioningItem(middle_tensor, safe_middle_frame, float(middle_image_weight)))
209
- if end_image_filepath:
210
- end_tensor = self._prepare_conditioning_tensor_from_file(end_image_filepath, height, width, padding_values)
211
- last_frame_index = actual_num_frames - 1
212
- conditioning_items.append(ConditioningItem(end_tensor, last_frame_index, float(end_image_weight)))
213
 
214
- # --- <LÓGICA CORRIGIDA E SIMPLIFICADA> ---
215
- elif mode == "video-to-video":
216
- indices_to_extract = list(range(0, int(frames_to_use), 8))
217
- extracted_frames_np = self._extract_frames_from_video(input_video_filepath, indices_to_extract)
218
- x=1
219
- with tempfile.TemporaryDirectory() as temp_dir:
220
- for i, frame_np in enumerate(extracted_frames_np):
221
- x = x+1
222
- frame_index = indices_to_extract[i]
223
- temp_frame_path = os.path.join(temp_dir, f"frame_{frame_index}.png")
224
- imageio.imwrite(temp_frame_path, frame_np)
225
-
226
- # Reutiliza a função de processamento de imagem, como você sugeriu
227
- frame_tensor = self._prepare_conditioning_tensor_from_file(
228
- temp_frame_path, height, width, padding_values
229
- )
230
- conditioning_items.append(ConditioningItem(frame_tensor, ((x*8)-8)-1, 0.5))
231
- print(f"[INFO] {len(conditioning_items)} frames do vídeo foram processados como keyframes de condicionamento.")
232
 
233
- call_kwargs = {
234
- "prompt": prompt, "negative_prompt": negative_prompt, "height": height_padded, "width": width_padded,
235
- "num_frames": actual_num_frames, "frame_rate": int(FPS), "generator": generator, "output_type": "pt",
236
- "conditioning_items": conditioning_items if conditioning_items else None,
237
- "media_items": None,
238
- "decode_timestep": self.config["decode_timestep"], "decode_noise_scale": self.config["decode_noise_scale"],
239
- "stochastic_sampling": True, #self.config["stochastic_sampling"], "image_cond_noise_scale": 0.15,
240
- "is_video": False, "vae_per_channel_normalize": True,
241
- "mixed_precision": True, #(self.config["precision"] == "mixed_precision"),
242
- "offload_to_cpu": False, "enhance_prompt": False,
243
- "skip_layer_strategy": None, #$/#SkipLayerStrategy.AttentionValues
244
- }
245
 
246
- result_tensor = None
247
- if improve_texture:
248
- if not self.latent_upsampler:
249
- raise ValueError("Upscaler espacial não carregado.")
250
- multi_scale_pipeline = LTXMultiScalePipeline(self.pipeline, self.latent_upsampler)
251
- first_pass_args = self.config.get("first_pass", {}).copy()
252
- first_pass_args["guidance_scale"] = float(guidance_scale)
253
- second_pass_args = self.config.get("second_pass", {}).copy()
254
- second_pass_args["guidance_scale"] = float(guidance_scale)
255
- multi_scale_call_kwargs = call_kwargs.copy()
256
- multi_scale_call_kwargs.update({"downscale_factor": self.config["downscale_factor"], "first_pass": first_pass_args, "second_pass": second_pass_args})
257
- result_tensor = multi_scale_pipeline(**multi_scale_call_kwargs).images
258
- log_tensor_info(result_tensor, "Resultado da Etapa 2 (Saída do Pipeline Multi-Scale)")
259
- else:
260
- single_pass_kwargs = call_kwargs.copy()
261
- first_pass_config = self.config.get("first_pass", {})
262
- single_pass_kwargs.update({
263
- "guidance_scale": float(guidance_scale),
264
- "stg_scale": first_pass_config.get("stg_scale"),
265
- "rescaling_scale": first_pass_config.get("rescaling_scale"),
266
- "skip_block_list": first_pass_config.get("skip_block_list"),
267
- "timesteps": first_pass_config.get("timesteps"),
268
- })
269
 
270
- print("\n[INFO] Executando pipeline de etapa única...")
271
- result_tensor = self.pipeline(**single_pass_kwargs).images
272
-
273
- pad_left, pad_right, pad_top, pad_bottom = padding_values
274
- slice_h_end = -pad_bottom if pad_bottom > 0 else None
275
- slice_w_end = -pad_right if pad_right > 0 else None
276
-
277
- result_tensor = result_tensor[:, :, :actual_num_frames, pad_top:slice_h_end, pad_left:slice_w_end]
278
- log_tensor_info(result_tensor, "Tensor Final (Após Pós-processamento, Antes de Salvar)")
 
 
 
279
 
280
- video_np = (result_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).astype(np.uint8)
281
- temp_dir = tempfile.mkdtemp()
282
- output_video_path = os.path.join(temp_dir, f"output_{used_seed}.mp4")
 
 
283
 
284
- with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], codec='libx264', quality=8) as writer:
285
- total_frames = len(video_np)
286
- for i, frame in enumerate(video_np):
287
- writer.append_data(frame)
288
- if progress_callback:
289
- progress_callback(i + 1, total_frames)
290
-
291
- self._log_gpu_memory("Fim da Geração")
292
- return output_video_path, used_seed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
293
 
294
- print("Criando instância do VideoService. O carregamento do modelo começará agora...")
295
  video_generation_service = VideoService()
 
1
  # video_service.py
2
 
 
3
  import torch
4
  import numpy as np
5
  import random
 
8
  from pathlib import Path
9
  import imageio
10
  import tempfile
 
11
  import sys
12
  import subprocess
13
+ import threading
14
+ import time
15
+ from huggingface_hub import hf_hub_download
16
 
17
+ # --- LÓGICA DE SETUP E DEPENDÊNCIAS ---
18
 
19
  def run_setup():
 
20
  setup_script_path = "setup.py"
21
  if not os.path.exists(setup_script_path):
22
  print("AVISO: script 'setup.py' não encontrado. Pulando a clonagem de dependências.")
 
35
  run_setup()
36
 
37
  def add_deps_to_path():
 
38
  if not LTX_VIDEO_REPO_DIR.exists():
39
  raise FileNotFoundError(f"Repositório LTX-Video não encontrado em '{LTX_VIDEO_REPO_DIR}'. Execute o setup.")
40
  if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
 
42
 
43
  add_deps_to_path()
44
 
45
+ # Importações específicas do modelo
46
  from inference import (
47
  create_ltx_video_pipeline, create_latent_upsampler,
48
  load_image_to_tensor_with_resize_and_crop, seed_everething,
49
  calculate_padding, load_media_file
50
  )
51
+ from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem
52
  from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
53
 
54
+ # --- CONFIGURAÇÃO DA DISTRIBUIÇÃO DE GPUS ---
55
+ GPU_MAPPING = [
56
+ {'base': 'cuda:0', 'upscaler': 'cuda:2'},
57
+ {'base': 'cuda:1', 'upscaler': 'cuda:3'}
58
+ ]
 
 
 
 
 
 
 
 
 
 
 
59
 
 
60
  class VideoService:
61
  def __init__(self):
62
+ print("Inicializando VideoService (modo Lazy Loading)...")
63
+ self.models_loaded = False
64
+ self.workers = None
65
  self.config = self._load_config()
66
+ self.models_dir = "downloaded_models"
67
+ self.loading_lock = threading.Lock() # Para evitar que múltiplos usuários iniciem o carregamento ao mesmo tempo
 
 
 
 
 
 
 
 
 
68
 
69
+ def _ensure_models_are_loaded(self):
70
+ """Verifica se os modelos estão carregados e os carrega se não estiverem."""
71
+ with self.loading_lock:
72
+ if not self.models_loaded:
73
+ print("Primeira requisição recebida. Iniciando carregamento dos modelos...")
74
+ if torch.cuda.is_available() and torch.cuda.device_count() < 4:
75
+ raise RuntimeError(f"Este serviço está configurado para 4 GPUs, mas apenas {torch.cuda.device_count()} foram encontradas.")
76
+
77
+ self._download_model_files()
78
+ self.workers = self._initialize_workers()
79
+ self.models_loaded = True
80
+ print(f"Modelos carregados com sucesso. {len(self.workers)} workers prontos.")
 
 
 
81
 
82
  def _load_config(self):
83
  config_file_path = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled.yaml"
84
  with open(config_file_path, "r") as file:
85
  return yaml.safe_load(file)
86
 
87
+ def _download_model_files(self):
88
+ Path(self.models_dir).mkdir(parents=True, exist_ok=True)
 
89
  LTX_REPO = "Lightricks/LTX-Video"
90
+ print("Baixando arquivos de modelo (se necessário)...")
91
+ self.distilled_model_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["checkpoint_path"], local_dir=self.models_dir)
92
+ self.spatial_upscaler_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["spatial_upscaler_model_path"], local_dir=self.models_dir)
93
+ print("Download de modelos concluído.")
 
 
 
 
 
 
 
 
 
 
 
94
 
95
+ def _load_models_for_worker(self, base_device, upscaler_device):
96
+ print(f"Carregando modelo base para {base_device} e upscaler para {upscaler_device}")
97
+ pipeline = create_ltx_video_pipeline(
98
+ ckpt_path=self.distilled_model_path, precision=self.config["precision"],
99
+ text_encoder_model_name_or_path=self.config["text_encoder_model_name_or_path"],
100
+ sampler=self.config["sampler"], device="cpu", enhance_prompt=False,
101
+ prompt_enhancer_image_caption_model_name_or_path=self.config["prompt_enhancer_image_caption_model_name_or_path"],
102
+ prompt_enhancer_llm_model_name_or_path=self.config["prompt_enhancer_llm_model_name_or_path"],
103
+ )
104
+ latent_upsampler = create_latent_upsampler(self.spatial_upscaler_path, device="cpu")
105
+ pipeline.to(base_device)
106
+ latent_upsampler.to(upscaler_device)
107
+ return pipeline, latent_upsampler
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
108
 
109
+ def _initialize_workers(self):
110
+ workers = []
111
+ for i, mapping in enumerate(GPU_MAPPING):
112
+ print(f"--- Inicializando Worker {i} ---")
113
+ pipeline, latent_upsampler = self._load_models_for_worker(mapping['base'], mapping['upscaler'])
114
+ workers.append({"id": i, "base_pipeline": pipeline, "latent_upsampler": latent_upsampler, "devices": mapping, "lock": threading.Lock()})
115
+ return workers
116
 
117
+ def _acquire_worker(self):
118
+ while True:
119
+ for worker in self.workers:
120
+ if worker["lock"].acquire(blocking=False):
121
+ print(f"Worker {worker['id']} adquirido para uma nova tarefa.")
122
+ return worker
123
+ time.sleep(0.1)
124
 
125
+ def generate(self, prompt, negative_prompt, input_image_filepath=None, input_video_filepath=None,
126
+ height=512, width=704, mode="text-to-video", duration=2.0,
127
+ frames_to_use=9, seed=42, randomize_seed=True, guidance_scale=1.0, # Ignorado, mas mantido por compatibilidade
128
+ improve_texture=True, progress_callback=None):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
129
 
130
+ # A MÁGICA DO LAZY LOADING ACONTECE AQUI
131
+ self._ensure_models_are_loaded()
132
 
133
+ worker = self._acquire_worker()
134
+ base_device = worker['devices']['base']
135
+ upscaler_device = worker['devices']['upscaler']
 
 
 
 
 
 
 
 
136
 
137
+ try:
138
+ # ... (todo o resto do código da função generate permanece exatamente o mesmo) ...
139
+ if mode == "image-to-video" and not input_image_filepath: raise ValueError("Caminho da imagem é obrigatório para o modo image-to-video")
140
+ if mode == "video-to-video" and not input_video_filepath: raise ValueError("Caminho do vídeo é obrigatório para o modo video-to-video")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
141
 
142
+ used_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
143
+ seed_everething(used_seed)
 
 
 
 
 
 
 
 
 
 
144
 
145
+ FPS = 24.0; MAX_NUM_FRAMES = 257
146
+ target_frames_rounded = round(duration * FPS)
147
+ n_val = round((float(target_frames_rounded) - 1.0) / 8.0)
148
+ actual_num_frames = max(9, min(MAX_NUM_FRAMES, int(n_val * 8 + 1)))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
149
 
150
+ height_padded = ((height - 1) // 32 + 1) * 32
151
+ width_padded = ((width - 1) // 32 + 1) * 32
152
+ padding_values = calculate_padding(height, width, height_padded, width_padded)
153
+ pad_left, pad_right, pad_top, pad_bottom = padding_values
154
+
155
+ call_kwargs_base = {
156
+ "prompt": prompt, "negative_prompt": negative_prompt, "num_frames": actual_num_frames, "frame_rate": int(FPS),
157
+ "decode_timestep": 0.05, "decode_noise_scale": self.config["decode_noise_scale"],
158
+ "stochastic_sampling": self.config["stochastic_sampling"], "image_cond_noise_scale": 0.025,
159
+ "is_video": True, "vae_per_channel_normalize": True, "mixed_precision": (self.config["precision"] == "mixed_precision"),
160
+ "offload_to_cpu": False, "enhance_prompt": False, "skip_layer_strategy": SkipLayerStrategy.AttentionValues
161
+ }
162
 
163
+ result_tensor = None
164
+ if improve_texture:
165
+ downscale_factor = self.config.get("downscale_factor", 0.5)
166
+ downscaled_height_ideal = int(height_padded * downscale_factor); downscaled_width_ideal = int(width_padded * downscale_factor)
167
+ downscaled_height = ((downscaled_height_ideal - 1) // 32 + 1) * 32; downscaled_width = ((downscaled_width_ideal - 1) // 32 + 1) * 32
168
 
169
+ # --- PASSE 1 ---
170
+ first_pass_kwargs = call_kwargs_base.copy()
171
+ first_pass_kwargs.update({
172
+ "height": downscaled_height, "width": downscaled_width,
173
+ "generator": torch.Generator(device=base_device).manual_seed(used_seed),
174
+ "output_type": "latent", "guidance_scale": 1.0,
175
+ "timesteps": self.config["first_pass"]["timesteps"],
176
+ "stg_scale": self.config["first_pass"]["stg_scale"],
177
+ "rescaling_scale": self.config["first_pass"]["rescaling_scale"],
178
+ "skip_block_list": self.config["first_pass"]["skip_block_list"]
179
+ })
180
+
181
+ if mode == "image-to-video":
182
+ padding_low_res = calculate_padding(downscaled_height, downscaled_width, downscaled_height, downscaled_width)
183
+ media_tensor_low_res = load_image_to_tensor_with_resize_and_crop(input_image_filepath, downscaled_height, downscaled_width)
184
+ media_tensor_low_res = torch.nn.functional.pad(media_tensor_low_res, padding_low_res)
185
+ first_pass_kwargs["conditioning_items"] = [ConditioningItem(media_tensor_low_res.to(base_device), 0, 1.0)]
186
+
187
+ print(f"Worker {worker['id']}: Iniciando passe 1 em {base_device}")
188
+ with torch.no_grad(): low_res_latents = worker['base_pipeline'](**first_pass_kwargs).images
189
+
190
+ low_res_latents = low_res_latents.to(upscaler_device)
191
+ with torch.no_grad(): high_res_latents = worker['latent_upsampler'](low_res_latents)
192
+ high_res_latents = high_res_latents.to(base_device)
193
+
194
+ # --- PASSE 2 ---
195
+ second_pass_kwargs = call_kwargs_base.copy()
196
+ high_res_h, high_res_w = downscaled_height * 2, downscaled_width * 2
197
+ second_pass_kwargs.update({
198
+ "height": high_res_h, "width": high_res_w, "latents": high_res_latents,
199
+ "generator": torch.Generator(device=base_device).manual_seed(used_seed),
200
+ "output_type": "pt", "image_cond_noise_scale": 0.0, "guidance_scale": 1.0,
201
+ "timesteps": self.config["second_pass"]["timesteps"],
202
+ "stg_scale": self.config["second_pass"]["stg_scale"],
203
+ "rescaling_scale": self.config["second_pass"]["rescaling_scale"],
204
+ "skip_block_list": self.config["second_pass"]["skip_block_list"],
205
+ "tone_map_compression_ratio": self.config["second_pass"].get("tone_map_compression_ratio", 0.0)
206
+ })
207
+
208
+ if mode == "image-to-video":
209
+ padding_high_res = calculate_padding(high_res_h, high_res_w, high_res_h, high_res_w)
210
+ media_tensor_high_res = load_image_to_tensor_with_resize_and_crop(input_image_filepath, high_res_h, high_res_w)
211
+ media_tensor_high_res = torch.nn.functional.pad(media_tensor_high_res, padding_high_res)
212
+ second_pass_kwargs["conditioning_items"] = [ConditioningItem(media_tensor_high_res.to(base_device), 0, 1.0)]
213
+
214
+ print(f"Worker {worker['id']}: Iniciando passe 2 em {base_device}")
215
+ with torch.no_grad(): result_tensor = worker['base_pipeline'](**second_pass_kwargs).images
216
+
217
+ else: # Passe Único
218
+ single_pass_kwargs = call_kwargs_base.copy()
219
+ first_pass_config = self.config["first_pass"]
220
+ single_pass_kwargs.update({
221
+ "height": height_padded, "width": width_padded, "output_type": "pt",
222
+ "generator": torch.Generator(device=base_device).manual_seed(used_seed),
223
+ "guidance_scale": 1.0, **first_pass_config
224
+ })
225
+ if mode == "image-to-video":
226
+ media_tensor_final = load_image_to_tensor_with_resize_and_crop(input_image_filepath, height_padded, width_padded)
227
+ media_tensor_final = torch.nn.functional.pad(media_tensor_final, padding_values)
228
+ single_pass_kwargs["conditioning_items"] = [ConditioningItem(media_tensor_final.to(base_device), 0, 1.0)]
229
+ elif mode == "video-to-video":
230
+ 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)
231
+
232
+ print(f"Worker {worker['id']}: Iniciando passe único em {base_device}")
233
+ with torch.no_grad(): result_tensor = worker['base_pipeline'](**single_pass_kwargs).images
234
+
235
+ if result_tensor.shape[-2:] != (height, width):
236
+ num_frames_final = result_tensor.shape[2]
237
+ videos_tensor = result_tensor.permute(0, 2, 1, 3, 4).reshape(-1, result_tensor.shape[1], result_tensor.shape[3], result_tensor.shape[4])
238
+ videos_resized = torch.nn.functional.interpolate(videos_tensor, size=(height, width), mode='bilinear', align_corners=False)
239
+ result_tensor = videos_resized.reshape(result_tensor.shape[0], num_frames_final, result_tensor.shape[1], height, width).permute(0, 2, 1, 3, 4)
240
+
241
+ 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)]
242
+ video_np = (result_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).astype(np.uint8)
243
+ temp_dir = tempfile.mkdtemp()
244
+ output_video_path = os.path.join(temp_dir, f"output_{used_seed}.mp4")
245
+
246
+ with imageio.get_writer(output_video_path, fps=call_kwargs_base["frame_rate"], codec='libx264', quality=8) as writer:
247
+ for i, frame in enumerate(video_np):
248
+ writer.append_data(frame)
249
+ if progress_callback: progress_callback(i + 1, len(video_np))
250
+ return output_video_path, used_seed
251
+
252
+ except Exception as e:
253
+ print(f"!!!!!!!! ERRO no Worker {worker['id']} !!!!!!!!\n{e}")
254
+ raise e
255
+ finally:
256
+ print(f"Worker {worker['id']}: Tarefa finalizada. Limpando cache e liberando worker...")
257
+ with torch.cuda.device(base_device): torch.cuda.empty_cache()
258
+ with torch.cuda.device(upscaler_device): torch.cuda.empty_cache()
259
+ worker["lock"].release()
260
 
261
+ # A instância do serviço é criada aqui, mas os modelos só serão carregados no primeiro clique.
262
  video_generation_service = VideoService()