Update dreamo_helpers.py

dreamo_helpers.py

@@ -1,288 +1,106 @@
-# --- app.py (O Painel de Controle do Maestro - Gestão de Memória Restaurada) ---
-# By Carlex & Gemini & DreamO
+# dreamo_helpers.py
+# Módulo de serviço OTIMIZADO para o DreamO. Sem redimensionamento automático.
 
-# --- Ato 1: A Convocação da Orquestra (Importações) ---
-import gradio as gr
-import torch
 import os
-import yaml
-from PIL import Image
-import shutil
-import gc
-import subprocess
-import math
-import google.generativeai as genai
+import cv2
+import torch
 import numpy as np
-import imageio
-from pathlib import Path
+from PIL import Image
 import huggingface_hub
-import json
-
-# --- Músicos (Implementações Externas) ---
-from inference import create_ltx_video_pipeline, load_image_to_tensor_with_resize_and_crop, seed_everething, calculate_padding
-from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem
-from dreamo_helpers import dreamo_generator_singleton
-
-# --- Ato 2: A Preparação do Palco (Configurações) ---
-config_file_path = "configs/ltxv-13b-0.9.8-distilled.yaml"
-with open(config_file_path, "r") as file:
-    PIPELINE_CONFIG_YAML = yaml.safe_load(file)
-
-LTX_REPO = "Lightricks/LTX-Video"
-models_dir = "downloaded_models_gradio_cpu_init"
-Path(models_dir).mkdir(parents=True, exist_ok=True)
-WORKSPACE_DIR = "aduc_workspace"
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-
-VIDEO_FPS = 30
-VIDEO_DURATION_SECONDS = 4
-VIDEO_TOTAL_FRAMES = VIDEO_DURATION_SECONDS * VIDEO_FPS
-
-# <<<< INICIALIZAÇÃO NA CPU >>>>
-print("Baixando e criando pipelines LTX na CPU...")
-distilled_model_actual_path = huggingface_hub.hf_hub_download(repo_id=LTX_REPO, filename=PIPELINE_CONFIG_YAML["checkpoint_path"], local_dir=models_dir, local_dir_use_symlinks=False)
-pipeline_instance = create_ltx_video_pipeline(ckpt_path=distilled_model_actual_path, precision=PIPELINE_CONFIG_YAML["precision"], text_encoder_model_name_or_path=PIPELINE_CONFIG_YAML["text_encoder_model_name_or_path"], sampler=PIPELINE_CONFIG_YAML["sampler"], device='cpu')
-print("Modelos LTX prontos (na CPU).")
-
-# --- Ato 3: As Partituras dos Músicos (Funções) ---
-
-def get_storyboard_from_director(num_fragments: int, prompt: str, initial_image_path: str, progress=gr.Progress()):
-    progress(0.5, desc="[Diretor Gemini] Criando o storyboard...")
-    if not initial_image_path: raise gr.Error("Por favor, forneça uma imagem de referência inicial.")
-    if not GEMINI_API_KEY: raise gr.Error("Chave da API Gemini não configurada!")
-    genai.configure(api_key=GEMINI_API_KEY)
-    try:
-        with open("prompts/director_storyboard_v2.txt", "r", encoding="utf-8") as f: template = f.read()
-    except FileNotFoundError: raise gr.Error("'prompts/director_storyboard_v2.txt' não encontrado!")
-    director_prompt = template.format(user_prompt=prompt, num_fragments=int(num_fragments))
-    model = genai.GenerativeModel('gemini-2.0-flash')
-    img = Image.open(initial_image_path)
-    response = model.generate_content([director_prompt, img])
-    try:
-        cleaned_response = response.text.strip().replace("```json", "").replace("```", "")
-        storyboard_data = json.loads(cleaned_response)
-        return storyboard_data.get("storyboard", [])
-    except Exception as e:
-        raise gr.Error(f"O Diretor retornou uma resposta inesperada. Erro: {e}")
-
-def run_keyframe_generation(storyboard, ref_img_path_1, ref_img_path_2, ref_task_1, ref_task_2):
-    if not storyboard: raise gr.Error("Nenhum roteiro para gerar imagens-chave.")
-    if not ref_img_path_1: raise gr.Error("A Referência 1 é obrigatória.")
-    
-    keyframe_paths, log_history = [], ""
-    try:
-        dreamo_generator_singleton.to_gpu()
-        for i, prompt in enumerate(storyboard):
-            log_message = f"Pintando Cena {i+1}/{len(storyboard)} com DreamO..."
-            log_history += log_message + "\n"
-            yield {keyframe_log_output: gr.update(value=log_history)}
-            
-            output_path = os.path.join(WORKSPACE_DIR, f"keyframe_image_{i+1}.png")
-            
-            image = dreamo_generator_singleton.generate_image(
-                ref_image1_np=np.array(Image.open(ref_img_path_1).convert("RGB")) if ref_img_path_1 else None,
-                ref_image2_np=np.array(Image.open(ref_img_path_2).convert("RGB")) if ref_img_path_2 else None,
-                ref_task1=ref_task_1, ref_task2=ref_task_2,
-                prompt=prompt
-            )
-            image.save(output_path)
-            keyframe_paths.append(output_path)
-            
-            log_message = f"Cena {i+1} pintada."
-            log_history += log_message + "\n"
-            yield {keyframe_log_output: gr.update(value=log_history), keyframe_gallery_output: gr.update(value=keyframe_paths), keyframe_images_state: keyframe_paths}
-    finally:
-        dreamo_generator_singleton.to_cpu()
+import gc
+from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+from torchvision.transforms.functional import normalize
+from dreamo.dreamo_pipeline import DreamOPipeline
+from dreamo.utils import img2tensor, tensor2img
+from tools import BEN2
+
+class Generator:
+    def __init__(self):
+        self.cpu_device = torch.device('cpu')
+        self.gpu_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        print("Carregando modelos DreamO para a CPU...")
+        model_root = 'black-forest-labs/FLUX.1-dev'
+        self.dreamo_pipeline = DreamOPipeline.from_pretrained(model_root, torch_dtype=torch.bfloat16)
+        self.dreamo_pipeline.load_dreamo_model(self.cpu_device, use_turbo=True)
         
-    log_history += "\nPintura de todas as cenas concluída!"
-    yield {keyframe_log_output: gr.update(value=log_history)}
-
-def run_ltx_animation(current_fragment_index, motion_prompt, conditioning_items_data, seed, cfg, progress=gr.Progress()):
-    progress(0, desc=f"[Animador LTX] Interpolando Cena {current_fragment_index}...")
-    output_path = os.path.join(WORKSPACE_DIR, f"fragment_{current_fragment_index}.mp4")
-    target_device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-    try:
-        pipeline_instance.to(target_device)
-        conditioning_items = []
-        for (path, start_frame, strength) in conditioning_items_data:
-            tensor = load_image_to_tensor_with_resize_and_crop(path, VIDEO_HEIGHT, VIDEO_WIDTH)
-            conditioning_items.append(ConditioningItem(tensor.to(target_device), start_frame, strength))
-
-        num_frames_padded = 121
-        padded_h, padded_w = ((VIDEO_HEIGHT - 1) // 32 + 1) * 32, ((VIDEO_WIDTH - 1) // 32 + 1) * 32
-        padding_vals = calculate_padding(VIDEO_HEIGHT, VIDEO_WIDTH, padded_h, padded_w)
-
-        for cond_item in conditioning_items:
-            cond_item.media_item = torch.nn.functional.pad(cond_item.media_item, padding_vals)
-
-        timesteps = PIPELINE_CONFIG_YAML.get("first_pass", {}).get("timesteps")
-        kwargs = {
-            "prompt": motion_prompt, "negative_prompt": "blurry, distorted, bad quality, artifacts",
-            "height": padded_h, "width": padded_w, "num_frames": num_frames_padded, "frame_rate": VIDEO_FPS,
-            "generator": torch.Generator(device=target_device).manual_seed(int(seed) + current_fragment_index),
-            "output_type": "pt", "guidance_scale": float(cfg), "timesteps": timesteps,
-            "conditioning_items": conditioning_items, "vae_per_channel_normalize": True,
-            "decode_timestep": PIPELINE_CONFIG_YAML["decode_timestep"],
-            "decode_noise_scale": PIPELINE_CONFIG_YAML["decode_noise_scale"],
-            "stochastic_sampling": PIPELINE_CONFIG_YAML["stochastic_sampling"],
-            "image_cond_noise_scale": 0.15, "is_video": True,
-            "mixed_precision": (PIPELINE_CONFIG_YAML["precision"] == "mixed_precision"),
-            "offload_to_cpu": False, "enhance_prompt": False
-        }
+        self.bg_rm_model = BEN2.BEN_Base().to(self.cpu_device).eval()
+        huggingface_hub.hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='models')
+        self.bg_rm_model.loadcheckpoints('models/BEN2_Base.pth')
         
-        result_tensor = pipeline_instance(**kwargs).images
-        pad_l, pad_r, pad_t, pad_b = padding_vals; slice_h, slice_w = (-pad_b if pad_b > 0 else None), (-pad_r if pad_r > 0 else None)
-        cropped_tensor = result_tensor[:, :, :VIDEO_TOTAL_FRAMES, pad_t:slice_h, pad_l:slice_w]; 
-        video_np = (cropped_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).astype(np.uint8)
-        with imageio.get_writer(output_path, fps=VIDEO_FPS, codec='libx264', quality=8) as writer:
-            for i, frame in enumerate(video_np): progress(i / len(video_np), desc=f"Renderizando frame {i+1}/{len(video_np)}..."); writer.append_data(frame)
-        return output_path
-    finally:
-        pipeline_instance.to('cpu')
+        self.face_helper = FaceRestoreHelper(
+            upscale_factor=1, face_size=512, crop_ratio=(1, 1),
+            det_model='retinaface_resnet50', save_ext='png', device=self.cpu_device,
+        )
+        print("Modelos DreamO prontos (na CPU).")
+
+    def to_gpu(self):
+        print("Movendo modelos DreamO para a GPU...")
+        self.dreamo_pipeline.to(self.gpu_device)
+        self.bg_rm_model.to(self.gpu_device)
+        self.face_helper.device = self.gpu_device
+        if hasattr(self.face_helper, 'face_parse'): self.face_helper.face_parse.to(self.gpu_device)
+        if hasattr(self.face_helper, 'face_det'): self.face_helper.face_det.to(self.gpu_device)
+        print("Modelos DreamO na GPU.")
+
+    def to_cpu(self):
+        print("Descarregando modelos DreamO da GPU...")
+        self.dreamo_pipeline.to(self.cpu_device)
+        self.bg_rm_model.to(self.cpu_device)
+        self.face_helper.device = self.cpu_device
+        if hasattr(self.face_helper, 'face_det'): self.face_helper.face_det.to(self.cpu_device)
+        if hasattr(self.face_helper, 'face_parse'): self.face_helper.face_parse.to(self.cpu_device)
         gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-def run_full_video_production(storyboard, keyframe_image_paths, seed, cfg):
-    if not storyboard or len(storyboard) != len(keyframe_image_paths):
-        raise gr.Error("Roteiro e imagens-chave devem existir e ter o mesmo número de cenas.")
-    video_fragments, log_history = [], ""
-    num_scenes = len(storyboard)
-    end_frame_index = VIDEO_TOTAL_FRAMES - 1
-    for i in range(num_scenes):
-        motion_prompt = storyboard[i]
-        start_image_path = keyframe_image_paths[i]
-        log_message = f"Preparando Cena {i+1}/{num_scenes}..."
-        log_history += log_message + "\n"
-        yield {video_production_log_output: gr.update(value=log_history)}
-        conditioning_items_data = [(start_image_path, 0, 1.0)]
-        log_message = f"  -> Ponto de Partida: {os.path.basename(start_image_path)} no quadro 0."
-        log_history += log_message + "\n"
-        if i < num_scenes - 1:
-            end_image_path = keyframe_image_paths[i+1]
-            conditioning_items_data.append((end_image_path, end_frame_index, 1.0))
-            log_message = f"  -> Ponto de Chegada: {os.path.basename(end_image_path)} no quadro {end_frame_index}."
-            log_history += log_message + "\n"
-        else:
-            log_message = f"  -> Ponto de Chegada: Animação livre (última cena)."
-            log_history += log_message + "\n"
-        yield {video_production_log_output: gr.update(value=log_history)}
-        fragment_path = run_ltx_animation(i + 1, motion_prompt, conditioning_items_data, seed, cfg)
-        video_fragments.append(fragment_path)
-        log_message = f"Cena {i+1} concluída e salva em {os.path.basename(fragment_path)}."
-        log_history += log_message + "\n"
-        yield {video_production_log_output: gr.update(value=log_history), fragment_gallery_output: gr.update(value=video_fragments), fragment_list_state: video_fragments, final_fragments_display: gr.update(value=video_fragments)}
-    log_history += "\nProdução de todas as cenas de vídeo concluída!"
-    yield {video_production_log_output: gr.update(value=log_history)}
-
-def concatenate_masterpiece(fragment_paths: list, progress=gr.Progress()):
-    progress(0.5, desc="Montando a obra-prima final..."); list_file_path, final_output_path = os.path.join(WORKSPACE_DIR, "concat_list.txt"), os.path.join(WORKSPACE_DIR, "obra_prima_final.mp4")
-    with open(list_file_path, "w") as f:
-        for path in fragment_paths: f.write(f"file '{os.path.abspath(path)}'\n")
-    command = f"ffmpeg -y -f concat -safe 0 -i {list_file_path} -c copy {final_output_path}"
-    try:
-        subprocess.run(command, shell=True, check=True, capture_output=True, text=True); return final_output_path
-    except subprocess.CalledProcessError as e:
-        raise gr.Error(f"FFmpeg falhou ao unir os vídeos: {e.stderr}")
-
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# LTX Video - Storyboard em Vídeo (ADUC-SDR)\n*By Carlex & Gemini & DreamO*")
-    
-    storyboard_state = gr.State([])
-    keyframe_images_state = gr.State([])
-    fragment_list_state = gr.State([])
-
-    if os.path.exists(WORKSPACE_DIR): shutil.rmtree(WORKSPACE_DIR)
-    os.makedirs(WORKSPACE_DIR)
-
-    with gr.Tabs():
-        with gr.TabItem("ETAPA 1: O DIRETOR (Roteiro Visual)"):
-            with gr.Row():
-                with gr.Column():
-                    num_fragments_input = gr.Slider(2, 10, 4, step=1, label="Número de Cenas")
-                    prompt_input = gr.Textbox(label="Ideia Geral (Prompt)")
-                    image_input = gr.Image(type="filepath", label="Imagem de Referência Principal")
-                    director_button = gr.Button("▶️ 1. Gerar Roteiro Visual", variant="primary")
-                with gr.Column():
-                    storyboard_to_show = gr.JSON(label="Roteiro Gerado (para visualização)")
-        with gr.TabItem("ETAPA 2: O PINTOR (Imagens-Chave)"):
-            with gr.Row():
-                with gr.Column(scale=2):
-                    gr.Markdown("### Controles do Pintor (DreamO)")
-                    with gr.Row():
-                        ref_image_1_input = gr.Image(label="Referência 1 (Principal)", type="filepath")
-                        ref_image_2_input = gr.Image(label="Referência 2 (Opcional, para composição)", type="filepath")
-                    with gr.Row():
-                        ref_task_1_input = gr.Dropdown(choices=["ip", "id", "style"], value="ip", label="Tarefa para Referência 1")
-                        ref_task_2_input = gr.Dropdown(choices=["ip", "id", "style"], value="ip", label="Tarefa para Referência 2")
-                    photographer_button = gr.Button("▶️ 2. Pintar Imagens-Chave", variant="primary")
-                    keyframe_log_output = gr.Textbox(label="Diário de Bordo do Pintor", lines=5, interactive=False)
-                with gr.Column(scale=1):
-                    keyframe_gallery_output = gr.Gallery(label="Imagens-Chave Pintadas", object_fit="contain", height="auto", type="filepath")
-        with gr.TabItem("ETAPA 3: A PRODUÇÃO (Gerar Cenas em Vídeo)"):
-            gr.Markdown(f"Gere o vídeo interpolando entre as imagens-chave. Cada clipe terá **{VIDEO_DURATION_SECONDS} segundos a {VIDEO_FPS} FPS ({VIDEO_TOTAL_FRAMES} quadros)**.")
-            with gr.Row():
-                with gr.Column():
-                    keyframes_to_render = gr.Gallery(label="Imagens-Chave para Animar", object_fit="contain", height="auto", interactive=False)
-                    animator_button = gr.Button("▶️ 3. Produzir Cenas em Vídeo", variant="primary")
-                    video_production_log_output = gr.Textbox(label="Diário de Bordo da Produção", lines=10, interactive=False)
-                with gr.Column():
-                    fragment_gallery_output = gr.Gallery(label="Cenas Produzidas (Vídeos)", object_fit="contain", height="auto")
-            with gr.Row():
-                seed_number = gr.Number(42, label="Seed")
-                cfg_slider = gr.Slider(1.0, 10.0, 2.5, step=0.1, label="CFG")
-        with gr.TabItem("ETAPA 4: PÓS-PRODUÇÃO"):
-            with gr.Row():
-                with gr.Column():
-                    editor_button = gr.Button("▶️ 4. Concatenar Vídeo Final", variant="primary")
-                    final_fragments_display = gr.JSON(label="Fragmentos a Concatenar")
-                with gr.Column():
-                    final_video_output = gr.Video(label="A Obra-Prima Final")
-
-    def director_success(storyboard_list, img_path):
-        if not storyboard_list: raise gr.Error("O storyboard está vazio ou em formato inválido.")
-        return {
-            storyboard_state: storyboard_list,
-            storyboard_to_show: gr.update(value=storyboard_list),
-            ref_image_1_input: gr.update(value=img_path)
-        }
-
-    director_button.click(
-        fn=get_storyboard_from_director,
-        inputs=[num_fragments_input, prompt_input, image_input],
-        outputs=[storyboard_state]
-    ).then(
-        fn=director_success,
-        inputs=[storyboard_state, image_input],
-        outputs=[storyboard_state, storyboard_to_show, ref_image_1_input]
-    )
-    
-    photographer_button.click(
-        fn=run_keyframe_generation,
-        inputs=[storyboard_state, ref_image_1_input, ref_image_2_input, ref_task_1_input, ref_task_2_input],
-        outputs=[keyframe_log_output, keyframe_gallery_output, keyframe_images_state]
-    ).then(
-        lambda paths: gr.update(value=paths),
-        inputs=[keyframe_images_state],
-        outputs=[keyframes_to_render]
-    )
-    
-    animator_button.click(
-        fn=run_full_video_production,
-        inputs=[storyboard_state, keyframe_images_state, seed_number, cfg_slider],
-        outputs=[video_production_log_output, fragment_gallery_output, fragment_list_state]
-    ).then(
-        lambda paths: gr.update(value=paths),
-        inputs=[fragment_list_state],
-        outputs=[final_fragments_display]
-    )
-
-    editor_button.click(
-        fn=concatenate_masterpiece,
-        inputs=[fragment_list_state],
-        outputs=[final_video_output]
-    )
-
-if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", share=True)
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+        print("GPU limpa após uso do DreamO.")
+
+    @torch.inference_mode()
+    def generate_image_with_gpu_management(self, ref_image1_np, ref_image2_np, ref_task1, ref_task2, prompt):
+        ref_conds = []
+        ref_images = [ref_image1_np, ref_image2_np]
+        ref_tasks = [ref_task1, ref_task2]
+        
+        for idx, (ref_image, ref_task) in enumerate(zip(ref_images, ref_tasks)):
+            if ref_image is not None:
+                # LÓGICA DE REDIMENSIONAMENTO REMOVIDA
+                if ref_task == "id":
+                    ref_image = self.get_align_face(ref_image)
+                elif ref_task != "style":
+                    ref_image = self.bg_rm_model.inference(Image.fromarray(ref_image))
+                
+                ref_image_tensor = img2tensor(np.array(ref_image), bgr2rgb=False).unsqueeze(0) / 255.0
+                ref_image_tensor = (2 * ref_image_tensor - 1.0).to(self.gpu_device, dtype=torch.bfloat16)
+                ref_conds.append({'img': ref_image_tensor, 'task': ref_task, 'idx': idx + 1})
+        
+        image = self.dreamo_pipeline(
+            prompt=prompt, width=1024, height=1024,
+            num_inference_steps=12, guidance_scale=4.5,
+            ref_conds=ref_conds, generator=torch.Generator(device="cpu").manual_seed(42)
+        ).images[0]
+        return image
+
+    @torch.no_grad()
+    def get_align_face(self, img):
+        self.face_helper.clean_all()
+        image_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        self.face_helper.read_image(image_bgr)
+        self.face_helper.get_face_landmarks_5(only_center_face=True)
+        self.face_helper.align_warp_face()
+        if len(self.face_helper.cropped_faces) == 0: return None
+        align_face = self.face_helper.cropped_faces[0]
+        input_tensor = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0
+        input_tensor = input_tensor.to(self.gpu_device)
+        parsing_out = self.face_helper.face_parse(normalize(input_tensor, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0]
+        parsing_out = parsing_out.argmax(dim=1, keepdim=True)
+        bg_label = [0, 16, 18, 7, 8, 9, 14, 15]
+        bg = sum(parsing_out == i for i in bg_label).bool()
+        white_image = torch.ones_like(input_tensor)
+        face_features_image = torch.where(bg, white_image, input_tensor)
+        return tensor2img(face_features_image, rgb2bgr=False)
+
+# --- Instância Singleton ---
+print("Inicializando o Pintor de Cenas (DreamO Helper)...")
+hf_token = os.getenv('HF_TOKEN')
+if hf_token: huggingface_hub.login(token=hf_token)
+dreamo_generator_singleton = Generator()
+print("Pintor de Cenas (DreamO Helper) pronto.")