app_wan.py

# app_wan.py

import os

# PyTorch 2.8 (temporary hack) — manter conforme app original
#os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')

import gradio as gr
import tempfile
import numpy as np
from PIL import Image
from gradio_client import Client, handle_file

# === Constantes (espelhando o app original) ===
MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"

# Dimensões
MAX_DIMENSION = 832
MIN_DIMENSION = 480
DIMENSION_MULTIPLE = 16
SQUARE_SIZE = 480

# Geração
MAX_SEED = np.iinfo(np.int32).max
FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81

MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)

default_negative_prompt = (
    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，"
    "JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，"
    "手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走,过曝，"
)

# === Importa o serviço de geração (manager) ===
from aduc_framework.managers.wan_manager import WanManager
wan_manager = WanManager()

# === Utilidades de UI ===
def switch_to_upload_tab():
    # Atualiza o Tabs existente para a aba "Upload"
    return gr.Tabs.update(selected="upload_tab")

def generate_end_frame(start_img, gen_prompt, progress=gr.Progress(track_tqdm=True)):
    """
    Chama uma API Gradio externa para gerar uma imagem (end frame).
    """
    if start_img is None:
        raise gr.Error("Please provide a Start Frame first.")

    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise gr.Error("HF_TOKEN not found in environment variables. Please set it in your Space secrets.")

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
        start_img.save(tmpfile.name)
        tmp_path = tmpfile.name

    progress(0.1, desc="Connecting to image generation API...")
    client = Client("multimodalart/nano-banana")

    progress(0.5, desc=f"Generating with prompt: '{gen_prompt}'...")
    try:
        result = client.predict(
            prompt=gen_prompt,
            images=[{"image": handle_file(tmp_path)}],
            manual_token=hf_token,
            api_name="/unified_image_generator",
        )
    finally:
        try:
            os.remove(tmp_path)
        except:
            pass

    progress(1.0, desc="Done!")
    return result

# Wrapper: a UI monta images_condition_items e delega ao serviço
def ui_generate_video(
    start_image_pil,
    start_frame_text,
    handle_image_pil,
    handle_frame_text,
    handle_peso,
    end_image_pil,
    end_frame_text,
    end_peso,
    prompt,
    negative_prompt=default_negative_prompt,
    duration_seconds=2.1,
    steps=8,
    guidance_scale=1.0,
    guidance_scale_2=1.0,
    seed=42,
    randomize_seed=False,
    progress=gr.Progress(track_tqdm=True),
):
    # Conversões simples
    def to_int_safe(v, default=0):
        try:
            return int(v)
        except:
            return default

    def to_float_safe(v, default=1.0):
        try:
            return float(v)
        except:
            return default

    # Peso da imagem inicial é fixo 1.0
    start_item = [start_image_pil, to_int_safe(start_frame_text, 0), 1.0]

    # Handle é opcional: se não houver imagem, não enviar item
    items = [start_item]
    if handle_image_pil is not None:
        items.append([handle_image_pil, to_int_safe(handle_frame_text, 4), to_float_safe(handle_peso, 1.0)])

    # Último item (end) sempre presente
    items.append([end_image_pil, to_int_safe(end_frame_text, MAX_FRAMES_MODEL - 1), to_float_safe(end_peso, 1.0)])

    return wan_manager.generate_video_from_conditions(
        images_condition_items=items,
        prompt=prompt,
        negative_prompt=negative_prompt,
        duration_seconds=float(duration_seconds),
        steps=int(steps),
        guidance_scale=float(guidance_scale),
        guidance_scale_2=float(guidance_scale_2),
        seed=int(seed),
        randomize_seed=bool(randomize_seed),
    )

# === UI Gradio ===
css = '''
.fillable{max-width: 1100px !important}
.dark .progress-text {color: white}
#general_items{margin-top: 2em}
#group_all{overflow:visible}
#group_all .styler{overflow:visible}
#group_tabs .tabitem{padding: 0}
.tab-wrapper{margin-top: -33px;z-index: 999;position: absolute;width: 100%;background-color: var(--block-background-fill);padding: 0;}
#component-9-button{width: 50%;justify-content: center}
#component-11-button{width: 50%;justify-content: center}
#or_item{text-align: center; padding-top: 1em; padding-bottom: 1em; font-size: 1.1em;margin-left: .5em;margin-right: .5em;width: calc(100% - 1em)}
#fivesec{margin-top: 5em;margin-left: .5em;margin-right: .5em;width: calc(100% - 1em)}
'''

with gr.Blocks(theme=gr.themes.Citrus(), css=css) as app:
    gr.Markdown("# Wan 2.2 Aduca-sdr")

    with gr.Row(elem_id="general_items"):
        with gr.Column():
            with gr.Group(elem_id="group_all"):
                with gr.Row():
                    # Coluna: Start (peso fixo 1.0)
                    with gr.Column():
                        start_image = gr.Image(type="pil", label="Start Frame", sources=["upload", "clipboard"])
                        start_frame_tb = gr.Textbox(label="Start Frame (int)", value="0")

                    # Coluna: Handle (opcional)
                    with gr.Column():
                        handle_image = gr.Image(type="pil", label="Handle Image", sources=["upload", "clipboard"])
                        handle_frame_tb = gr.Textbox(label="Handle Frame (int)", value="4")
                        handle_peso_sl = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="Handle Peso")

                    # Coluna: End (Upload/Generate)
                    with gr.Tabs(elem_id="group_tabs") as tabs:
                        with gr.TabItem("Upload", id="upload_tab"):
                            with gr.Column():
                                end_image = gr.Image(type="pil", label="End Frame", sources=["upload", "clipboard"])
                                end_frame_tb = gr.Textbox(label="End Frame (int)", value=str(MAX_FRAMES_MODEL - 1))
                                end_peso_sl = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="End Peso")
                        with gr.TabItem("Generate", id="generate_tab"):
                            generate_5seconds = gr.Button("Generate scene 5 seconds in the future", elem_id="fivesec")
                            gr.Markdown(
                                "Generate a custom end-frame with an edit model like Nano Banana or Qwen Image Edit",
                                elem_id="or_item",
                            )
                prompt = gr.Textbox(label="Prompt", info="Describe the transition between the two images")

                with gr.Accordion("Advanced Settings", open=False):
                    duration_seconds_input = gr.Slider(
                        minimum=MIN_DURATION,
                        maximum=MAX_DURATION,
                        step=0.1,
                        value=2.1,
                        label="Video Duration (seconds)",
                        info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.",
                    )
                    negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                    steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=8, label="Inference Steps")
                    guidance_scale_input = gr.Slider(
                        minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - high noise"
                    )
                    guidance_scale_2_input = gr.Slider(
                        minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - low noise"
                    )
                    with gr.Row():
                        seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
                        randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True)

                generate_button = gr.Button("Generate Video", variant="primary")

        with gr.Column():
            output_video = gr.Video(label="Generated Video", autoplay=True)

    # Inputs/outputs para o wrapper
    ui_inputs = [
        start_image, start_frame_tb,
        handle_image, handle_frame_tb, handle_peso_sl,
        end_image, end_frame_tb, end_peso_sl,
        prompt, negative_prompt_input, duration_seconds_input,
        steps_slider, guidance_scale_input, guidance_scale_2_input,
        seed_input, randomize_seed_checkbox,
    ]
    ui_outputs = [output_video, seed_input]

    generate_button.click(fn=ui_generate_video, inputs=ui_inputs, outputs=ui_outputs)

    # Cadeia “5 seconds”: alterna aba, gera end frame e dispara render
    generate_5seconds.click(
        fn=switch_to_upload_tab,
        inputs=None,
        outputs=[tabs]
    ).then(
        fn=lambda img: generate_end_frame(
            img,
            "this image is a still frame from a movie. generate a new frame with what happens on this scene 5 seconds in the future"
        ),
        inputs=[start_image],
        outputs=[end_image]
    ).success(
        fn=ui_generate_video,
        inputs=ui_inputs,
        outputs=ui_outputs
    )

if __name__ == "__main__":
    os.makedirs("examples", exist_ok=True)
    try:
        Image.new('RGB', (832, 480), color=(73, 109, 137)).save("examples/frame_1.png")
        Image.new('RGB', (832, 480), color=(173, 109, 237)).save("examples/frame_2.png")
        Image.new('RGB', (832, 480), color=(255, 255, 0)).save("examples/frame_3.png")
    except:
        pass

    app.launch(server_name="0.0.0.0", server_port=7860, show_error=True)