"""SOLRICKS 3D Animation Style for MiniMax H3 — ZeroGPU T2V and I2V demo."""
from __future__ import annotations
print('[boot] app.py started', flush=True)
import faulthandler
faulthandler.enable()
faulthandler.dump_traceback_later(180, repeat=True)
import os
import secrets
import tempfile
import traceback
from pathlib import Path
# Must precede torch/diffusers: ZeroGPU intercepts startup CUDA allocations.
print('[boot] importing ZeroGPU runtime', flush=True)
import spaces
print('[boot] ZeroGPU runtime imported; importing Gradio', flush=True)
import gradio as gr
print('[boot] Gradio imported', flush=True)
from h3_helpers import CANVASES, DEFAULT_CANVAS, FPS, frame_count, prepare_image, prepare_prompt

MODEL_REPO = 'MiniMaxAI/MiniMax-H3'
STYLE_REPO = 'SOLRICKS/3D-Animation-Style-MiniMax-H3'
STYLE_FILE = '3dsrx_minimax_h3_rank16-step00001250.safetensors'
TURBO_REPO = 'lightx2v/Minimax-h3-Turbo'
TURBO_FILE = 'minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors'
CONDITIONER = os.getenv('H3_CONDITIONER', 'multimodalart/qwen3vl-conditioner')
OUTPUT_DIR = Path(tempfile.gettempdir()) / 'solricks-3dsrx-outputs'
OUTPUT_DIR.mkdir(exist_ok=True)
PIPE = None
LOAD_ERROR = ''


def load_models():
    global PIPE, LOAD_ERROR
    try:
        print('[load 1/7] importing adapter tools', flush=True)
        from huggingface_hub import hf_hub_download
        from h3_adapters import read_adapter, install_adapter
        # Check gated access before downloading the large base.
        print('[load 2/7] downloading or opening 3dsrx style checkpoint', flush=True)
        style = read_adapter(hf_hub_download(STYLE_REPO, STYLE_FILE, token=os.getenv('HF_TOKEN')))
        print('[load 3/7] style adapter prepared; downloading or opening Turbo checkpoint', flush=True)
        turbo = read_adapter(hf_hub_download(TURBO_REPO, TURBO_FILE, token=False))
        print('[load 4/7] adapters prepared; importing H3 pipeline', flush=True)
        import torch
        from diffusers import ComponentsManager
        from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
        from h3_split_blocks import AnimationGeneratorBlocks
        MiniMaxH3ModularPipeline.min_duration = property(lambda self: 4.0)
        manager = ComponentsManager()
        pipe = AnimationGeneratorBlocks().init_pipeline(
            MODEL_REPO, components_manager=manager, collection='solricks-h3')
        print('[load 5/7] downloading and loading H3 transformer and VAEs (about 77 GB on a cold cache)', flush=True)
        pipe.load_components(dtype=torch.bfloat16)
        print('[3dsrx] installing Turbo adapter', flush=True)
        install_adapter(pipe.transformer, turbo, 'turbo')
        print('[3dsrx] installing 3D Animation Style adapter', flush=True)
        install_adapter(pipe.transformer, style, 'animation')
        pipe.transformer.set_adapters(['turbo', 'animation'], weights=[1.0, 0.8])
        pipe.transformer.set_attention_backend('_native_cudnn')
        # Pack only the transformer, keeping duplicated tensors under the disk quota.
        print('[load 6/7] preparing transformer for ZeroGPU', flush=True)
        pipe.transformer.to('cuda')
        PIPE = pipe
        print('[load 7/7] models ready', flush=True)
    except Exception as error:
        from huggingface_hub.errors import GatedRepoError
        traceback.print_exc()
        if isinstance(error, GatedRepoError):
            LOAD_ERROR = ('Model access is required. The Space owner needs to add an HF_TOKEN '
                          'secret with read access to the 3D Animation Style model, then restart the Space.')
        else:
            LOAD_ERROR = 'The generation model could not start. The Space owner can check the runtime logs.'


def gpu_duration(prompt_embeds, tags, image, height, width, frames, seed, strength):
    latent_frames = (frames - 5) // 17 * 5 + 2
    patches = (height // 32) * (width // 32)
    rows = latent_frames * patches + int(image is not None) * patches
    denoise_time = 8 * (1.1745e-4 * rows + 3.8396e-9 * rows ** 2)
    return max(90, int(denoise_time * 1.8 + 65))


@spaces.GPU(duration=gpu_duration, size='xlarge')
def denoise(prompt_embeds, tags, image, height, width, frames, seed, strength):
    import torch
    with torch.inference_mode():
        PIPE.vae.to('cuda')
        PIPE.audio_vae.to('cuda')
        PIPE.transformer.set_adapters(['turbo', 'animation'], weights=[1.0, float(strength)])
        result = PIPE(prompt_embeds=prompt_embeds.to('cuda'), text_token_tags=tags, image=image,
                      height=int(height), width=int(width), num_frames=int(frames), num_inference_steps=8,
                      generator=torch.Generator('cpu').manual_seed(int(seed)))
        return result.get('videos')[0], result.get('audio')[0].cpu(), result.get('sampling_rate')


def generate(mode, prompt, image, canvas, duration, strength, seed, randomize,
             request: gr.Request, progress=gr.Progress()):
    if LOAD_ERROR or PIPE is None:
        raise gr.Error(LOAD_ERROR or 'The generation model is still loading.')
    try:
        full_prompt = prepare_prompt(prompt)
        frames = frame_count(duration)
        if canvas not in CANVASES:
            raise ValueError('Choose one of the available video formats.')
        if not 0 <= float(strength) <= 1.2:
            raise ValueError('LoRA strength must be between 0 and 1.2.')
        seed = secrets.randbelow(2 ** 31) if randomize else int(seed)
        if not 0 <= seed < 2 ** 31:
            raise ValueError('Seed must be between 0 and 2147483647.')
        if mode not in ('Text to Video', 'Image to Video'):
            raise ValueError('Choose Text to Video or Image to Video.')
        if mode == 'Image to Video' and not image:
            raise ValueError('Upload a starting image for Image to Video.')
    except (ValueError, TypeError) as error:
        raise gr.Error(str(error)) from error
    from gradio_client import Client, handle_file
    from safetensors import safe_open
    from PIL import Image
    from diffusers.utils import encode_video
    # Keep the reference deployment's guard on image-conditioned requests.
    if mode == 'Image to Video':
        import ncii_guard
        if ncii_guard.classify(full_prompt)['label'] == 'ncii':
            raise gr.Error('This prompt was flagged by a content filter and was not run.')
    # Forward only the visitor's quota header, never the owner's HF_TOKEN.
    ip_token = request.headers.get('x-ip-token') if request is not None else None
    headers = {'x-ip-token': ip_token} if ip_token else {}
    try:
        with tempfile.TemporaryDirectory(prefix='3dsrx-input-') as directory:
            prepared = prepare_image(image, canvas, directory) if mode == 'Image to Video' else None
            progress(0.05, desc='Preparing your scene')
            client = Client(CONDITIONER, headers=headers, token=False)
            embedding_file, plan = client.predict(
                prompt=full_prompt, image_path=handle_file(prepared) if prepared else None,
                last_image_path=None, canvas=canvas, num_frames=frames,
                rewrite_prompt=False, api_name='/encode')
            with safe_open(embedding_file, framework='pt', device='cpu') as handle:
                meta = handle.metadata()
                embeds, tags = handle.get_tensor('prompt_embeds'), handle.get_tensor('text_token_tags')
            height, width, frames = (int(meta[key]) for key in ('height', 'width', 'num_frames'))
            first_frame = Image.open(prepared).convert('RGB') if prepared else None
            progress(0.25, desc='Generating 3D animation')
            video, audio, rate = denoise(embeds, tags, first_frame, height, width, frames, seed, strength)
            output = OUTPUT_DIR / f'3dsrx-{secrets.token_hex(12)}.mp4'
            progress(0.95, desc='Finishing your video')
            encode_video(video, fps=FPS, output_path=str(output), audio=audio, audio_sample_rate=rate)
        return str(output), f'{width} × {height} · {frames / FPS:.2f} s · LoRA {float(strength):.2f} · Seed {seed}', seed
    except gr.Error:
        raise
    except Exception as error:
        traceback.print_exc()
        if any(word in str(error).lower() for word in ('quota', 'gpu limit', 'allocate', 'capacity', 'concurrent', 'schedule')):
            raise gr.Error('ZeroGPU is busy or your GPU time is used up. Please try again later.') from error
        raise gr.Error('Generation could not finish. Please retry; the Space logs contain the error details.') from error


# Match the Generate Visuals palette in both browser color schemes.
SOLRICKS_COLORS = {
    'body_background_fill': '#0a1020',
    'body_text_color': '#f2f4f8',
    'body_text_color_subdued': '#b8c2d4',
    'background_fill_primary': '#0b1423',
    'background_fill_secondary': '#18243a',
    'border_color_primary': '#2c3b56',
    'border_color_accent': '#7268ff',
    'border_color_accent_subdued': '#3b456f',
    'color_accent_soft': '#252d50',
    'block_background_fill': 'transparent',
    'block_border_color': 'transparent',
    'block_info_text_color': '#b8c2d4',
    'block_label_background_fill': 'transparent',
    'block_label_border_color': 'transparent',
    'block_label_text_color': '#c4ccda',
    'block_title_background_fill': 'transparent',
    'block_title_border_color': 'transparent',
    'block_title_text_color': '#f2f4f8',
    'block_shadow': 'none',
    'panel_background_fill': '#101a2d',
    'panel_border_color': '#2c3b56',
    'input_background_fill': '#252d3d',
    'input_background_fill_focus': '#252d3d',
    'input_background_fill_hover': '#2b3549',
    'input_border_color': '#424b5e',
    'input_border_color_focus': '#8692ff',
    'input_border_color_hover': '#59657c',
    'input_placeholder_color': '#98a4b9',
    'input_shadow': 'none',
    'input_shadow_focus': '0 0 0 2px #7268ff30',
    'checkbox_background_color': '#252d3d',
    'checkbox_background_color_focus': '#343e58',
    'checkbox_background_color_hover': '#343e58',
    'checkbox_background_color_selected': '#7268ff',
    'checkbox_border_color': '#56617a',
    'checkbox_border_color_focus': '#8692ff',
    'checkbox_border_color_hover': '#8692ff',
    'checkbox_border_color_selected': '#7268ff',
    'checkbox_label_background_fill': '#252d3d',
    'checkbox_label_background_fill_hover': '#2b3549',
    'checkbox_label_background_fill_selected': '#2d365e',
    'checkbox_label_border_color': '#424b5e',
    'checkbox_label_border_color_hover': '#677393',
    'checkbox_label_border_color_selected': '#7d89f5',
    'checkbox_label_text_color': '#c4ccda',
    'checkbox_label_text_color_selected': '#ffffff',
    'checkbox_label_shadow': 'none',
    'accordion_text_color': '#f2f4f8',
    'button_primary_background_fill': 'linear-gradient(90deg, #7268ff, #596dff)',
    'button_primary_background_fill_hover': 'linear-gradient(90deg, #7a71ff, #6276ff)',
    'button_primary_border_color': '#7268ff',
    'button_primary_border_color_hover': '#9c9aff',
    'button_primary_text_color': '#ffffff',
    'button_primary_text_color_hover': '#ffffff',
    'button_primary_shadow': 'none',
    'button_secondary_background_fill': '#262d3c',
    'button_secondary_background_fill_hover': '#2d3547',
    'button_secondary_border_color': '#424b5e',
    'button_secondary_border_color_hover': '#677393',
    'button_secondary_text_color': '#f2f4f8',
    'button_secondary_text_color_hover': '#ffffff',
    'button_secondary_shadow': 'none',
    'link_text_color': '#a8b7ff',
    'link_text_color_active': '#d1d9ff',
    'link_text_color_hover': '#d1d9ff',
    'link_text_color_visited': '#a8b7ff',
    'code_background_fill': '#252d3d',
    'table_text_color': '#f2f4f8',
    'table_border_color': '#3a455a',
    'table_even_background_fill': '#18243a',
    'table_odd_background_fill': '#101a2d',
    'table_row_focus': '#2d365e',
    'loader_color': '#8190ff',
    'slider_color': '#7268ff',
}
SOLRICKS_THEME = gr.themes.Base(
    primary_hue='indigo', secondary_hue='blue', neutral_hue='slate',
    font=['Inter', 'ui-sans-serif', 'system-ui', 'sans-serif'],
).set(
    **SOLRICKS_COLORS,
    **{f'{key}_dark': value for key, value in SOLRICKS_COLORS.items()},
    color_accent='#7268ff', block_radius='0px', input_radius='0px',
    container_radius='0px', button_large_radius='0px',
    button_medium_radius='0px', button_small_radius='0px',
    block_label_text_weight='600', block_label_shadow='none',
    shadow_drop='none', shadow_drop_lg='none',
)

CSS = '''
html, body {margin:0;background:#0a1020!important;}
.gradio-container {
    max-width:none!important;min-height:100dvh!important;margin:0!important;
    padding:36px 24px 24px!important;color:#f2f4f8!important;
    background:
        radial-gradient(circle at 18% 12%,rgba(32,107,194,.22),transparent 28%),
        radial-gradient(circle at 82% 16%,rgba(78,54,177,.20),transparent 34%),
        linear-gradient(135deg,#08131f 0%,#0a1020 50%,#15173b 100%)!important;
}
.gradio-container .main,.gradio-container main {
    max-width:1160px!important;width:100%;margin-inline:auto!important;
}
#hero,#h3-input-panel,#h3-output-panel {
    border:1px solid rgba(125,167,255,.17)!important;border-radius:0!important;
    background:linear-gradient(90deg,rgba(28,41,70,.84) 0%,
        rgba(10,18,33,.95) 42%,rgba(6,13,25,.99) 100%)!important;
    box-shadow:none!important;
}
#hero {padding:28px 30px;margin-bottom:2px;}
#hero h1 {
    font-size:clamp(28px,4vw,40px);line-height:1.15;margin:12px 0;
    color:#f2f4f8;letter-spacing:-.035em;font-weight:800;
}
#hero p {font-size:15px;max-width:780px;line-height:1.65;color:#b8c2d4;}
#hero .eyebrow {font-size:11px;letter-spacing:.18em;font-weight:700;color:#a8b7ff;}
#hero .chips {display:flex;flex-wrap:wrap;gap:8px;margin-top:18px;}
#hero .chips span {
    border:1px solid rgba(125,167,255,.24);background:rgba(77,107,160,.16);
    color:#ced8ef;border-radius:3px;padding:5px 11px;font-size:12px;
}
#h3-main-grid {gap:18px!important;}
#h3-input-panel,#h3-output-panel {padding:24px!important;}
#h3-output-panel {background:linear-gradient(90deg,rgba(28,41,70,.74),rgba(6,13,25,.99))!important;}
#h3-input-panel .form,#h3-output-panel .form {
    background:transparent!important;border:0!important;box-shadow:none!important;
}
.gradio-container [data-testid="block-info"] {
    color:#c4ccda!important;background:transparent!important;
    font-size:12px!important;font-weight:600!important;letter-spacing:.035em;
}
#h3-prompt textarea,#h3-report textarea {
    background:#252d3d!important;color:#f2f4f8!important;
    border-radius:0!important;line-height:1.6!important;
}
#h3-prompt textarea {min-height:142px;padding:16px!important;}
#h3-report textarea {color:#c4ccda!important;}
#h3-video {background:rgba(4,10,20,.35)!important;border:1px solid #2c3b56!important;}
#generate {
    min-height:50px;background:linear-gradient(90deg,#7268ff,#596dff)!important;
    color:#fff!important;border:1px solid rgba(114,104,255,.22)!important;
    border-radius:0!important;font-weight:700!important;box-shadow:none!important;
}
#generate:hover {background:linear-gradient(90deg,#7a71ff,#6276ff)!important;}
#generate:focus-visible {outline:2px solid #b0b8ff;outline-offset:3px;}
.gradio-container .prose {color:#b8c2d4;}
.gradio-container .prose strong {color:#e5ebf5;}
.gradio-container .prose a {color:#a8b7ff!important;}
#model-status {font-size:13px;color:#b8c2d4;}
@media (max-width:700px) {
    .gradio-container {padding:20px 12px!important;}
    #hero {padding:24px 20px;}
    #h3-input-panel,#h3-output-panel {padding:18px!important;}
}
.gradio-container footer,
[aria-label="Gradio footer navigation"] {
    display: none !important;
}
'''


def change_mode(mode):
    is_image = mode == 'Image to Video'
    return gr.update(visible=is_image), 0.4 if is_image else 0.8


with gr.Blocks(title='3D Animation Style MiniMax H3', delete_cache=(3600, 3600)) as demo:
    gr.HTML('''<div id="hero"><div class="eyebrow">SOLRICKS / LoRA MODELS</div>
    <h1>Bring your imagination to life in 3D.</h1>
    <p>Create 3D animation from a scene description or a starting image
    with MiniMax H3.</p>
    <div class="chips"><span>Text to Video</span><span>Image to Video</span>
    <span>3dsrx</span><span>Native audio</span></div></div>''')
    with gr.Row(equal_height=True, elem_id='h3-main-grid'):
        with gr.Column(scale=1, elem_id='h3-input-panel'):
            mode = gr.Radio(['Text to Video', 'Image to Video'], value='Text to Video', label='Create with')
            image = gr.Image(type='filepath', label='Starting image', visible=False, height=280)
            prompt = gr.Textbox(label='Your scene', lines=5, elem_id='h3-prompt',
                value='A curious young girl in a mustard hoodie looks through a telescope on a rooftop at dusk. '
                      'She slowly looks up at the stars and smiles. Soft twilight, expressive character animation, one continuous shot.',
                placeholder='Describe the character, setting, action and camera movement…')
            gr.Markdown('The **3dsrx** style trigger is added automatically.')
            with gr.Row():
                canvas = gr.Dropdown(list(CANVASES), value=DEFAULT_CANVAS, label='Video format')
                duration = gr.Slider(4, 8, value=5, step=1, label='Duration (seconds)')
            with gr.Accordion('Style & Seed', open=False):
                strength = gr.Slider(0, 1.2, value=0.8, step=0.05, label='3D animation strength')
                gr.Markdown('Recommended: **0.8** for text, **0.4** for images. Try **1.0** for animals and fantasy scenes.')
                with gr.Row():
                    seed = gr.Number(value=42, precision=0, label='Seed', minimum=0, maximum=2147483647)
                    randomize = gr.Checkbox(value=True, label='Random seed')
            button = gr.Button('Generate animation', variant='primary', size='lg', elem_id='generate')
        with gr.Column(scale=1, elem_id='h3-output-panel'):
            output = gr.Video(label='Your animation', height=420, autoplay=False, format='mp4', elem_id='h3-video')
            report = gr.Textbox(label='Generation details', interactive=False, elem_id='h3-report')
            gr.Markdown('**Tips:** Start with one subject and a simple action. For images, choose a matching video format; '
                        'the starting frame is center-cropped to fit.')
    with gr.Accordion('Example scenes', open=False):
        gr.Examples([
            ['A playful raccoon and a tiny hedgehog explore a sunlit garden. The raccoon carefully offers the hedgehog a red berry. Gentle camera push, warm light, expressive 3D character animation.'],
            ['A cheerful otter floats on its back in a calm river, holding a red apple with both paws. It takes a small bite, then glances toward the camera. Soft sunlight, detailed fur, one continuous shot.'],
            ['A woman in a cozy cafe holds a ceramic cup with both hands. She looks out the window and smiles softly. Warm morning light, shallow depth of field, subtle natural movement.'],
        ], inputs=[prompt], cache_examples=False)
    gr.Markdown(f'[Model & downloads](https://huggingface.co/{STYLE_REPO}) · '
                'Prompts and starting images are sent to the shared '
                '[H3 conditioner](https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner) for scene encoding.')
    model_status = gr.Markdown(elem_id='model-status')
    mode.change(change_mode, [mode], [image, strength], queue=False)
    button.click(generate, [mode, prompt, image, canvas, duration, strength, seed, randomize],
                 [output, report, seed], concurrency_limit=1, api_name='generate')
    demo.load(lambda: LOAD_ERROR or '', outputs=model_status, queue=False)

print('[boot] interface constructed', flush=True)

if __name__ == '__main__':
    load_models()
    faulthandler.cancel_dump_traceback_later()
    print('[boot] launching Gradio server', flush=True)
    demo.queue(max_size=10, default_concurrency_limit=1).launch(
        theme=SOLRICKS_THEME, css=CSS,
        allowed_paths=[str(OUTPUT_DIR)], show_error=False)