# ginigen/Flux-VIDEO-BACKUP 소스참고

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from gradio_client import Client, handle_file
import random
import tempfile
import os
import logging
import torch
from diffusers import AutoencoderKL, TCDScheduler
from diffusers.models.model_loading_utils import load_state_dict
from huggingface_hub import hf_hub_download
from pathlib import Path
import torchaudio
from einops import rearrange
from scipy.io import wavfile
from transformers import pipeline
import cv2

# 환경 변수 설정으로 torch.load 체크 우회 (임시 해결책)
os.environ["TRANSFORMERS_ALLOW_UNSAFE_DESERIALIZATION"] = "1"

# Spaces GPU
try:
    import spaces
except:
    # GPU 데코레이터가 없을 때를 위한 더미 데코레이터
    class spaces:
        @staticmethod
        def GPU(duration=None):
            def decorator(func):
                return func
            return decorator

# API URLs
TEXT2IMG_API_URL = "http://211.233.58.201:7896"
VIDEO_API_URL = "http://211.233.58.201:7875"

# 로깅 설정
logging.basicConfig(level=logging.INFO)

# Image size presets
IMAGE_PRESETS = {
    "Custom": {"width": 1024, "height": 1024},
    "1:1 Square": {"width": 1024, "height": 1024},
    "4:3 Standard": {"width": 1024, "height": 768},
    "16:9 Widescreen": {"width": 1024, "height": 576},
    "9:16 Portrait": {"width": 576, "height": 1024},
    "6:19 Ultra Tall": {"width": 324, "height": 1024},
    "Instagram Square": {"width": 1080, "height": 1080},
    "Instagram Story": {"width": 1080, "height": 1920},
    "Instagram Landscape": {"width": 1080, "height": 566},
    "Facebook Cover": {"width": 820, "height": 312},
    "Twitter Header": {"width": 1500, "height": 500},
    "YouTube Thumbnail": {"width": 1280, "height": 720},
    "LinkedIn Banner": {"width": 1584, "height": 396},
}

def update_dimensions(preset):
    if preset in IMAGE_PRESETS:
        return IMAGE_PRESETS[preset]["width"], IMAGE_PRESETS[preset]["height"]
    return 1024, 1024

def add_watermark_to_video(input_video_path, output_video_path, watermark_text="Ginigen.com"):
    """Add watermark to video in the bottom-right corner"""
    try:
        # Open video
        cap = cv2.VideoCapture(input_video_path)
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # Video writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
        
        # Watermark settings
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = min(width, height) / 500  # Scale based on video size
        font_thickness = max(1, int(font_scale * 2))
        
        # Get text size
        (text_width, text_height), baseline = cv2.getTextSize(
            watermark_text, font, font_scale, font_thickness
        )
        
        # Position in bottom-right corner with padding
        padding = 20
        text_x = width - text_width - padding
        text_y = height - padding
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Add semi-transparent background for better visibility
            overlay = frame.copy()
            cv2.rectangle(overlay, 
                         (text_x - 10, text_y - text_height - 10),
                         (text_x + text_width + 10, text_y + 10),
                         (0, 0, 0), -1)
            frame = cv2.addWeighted(frame, 0.7, overlay, 0.3, 0)
            
            # Add watermark text
            cv2.putText(frame, watermark_text, (text_x, text_y), 
                       font, font_scale, (255, 255, 255), font_thickness, cv2.LINE_AA)
            
            out.write(frame)
        
        cap.release()
        out.release()
        cv2.destroyAllWindows()
        
        return True
    except Exception as e:
        logging.error(f"Watermark error: {str(e)}")
        return False

def generate_text_to_image(prompt, width, height, guidance, inference_steps, seed):
    if not prompt:
        return None, "Please enter a prompt"
    
    try:
        client = Client(TEXT2IMG_API_URL)
        if seed == -1:
            seed = random.randint(0, 9999999)
            
        result = client.predict(
            prompt=prompt,
            width=int(width),
            height=int(height),
            guidance=float(guidance),
            inference_steps=int(inference_steps),
            seed=int(seed),
            do_img2img=False,
            init_image=None,
            image2image_strength=0.8,
            resize_img=True,
            api_name="/generate_image"
        )
        return result[0], f"Seed used: {result[1]}"
    except Exception as e:
        logging.error(f"Image generation error: {str(e)}")
        return None, f"Error: {str(e)}"

def generate_video_from_image(image, prompt="", length=4.0):
    if image is None:
        return None
    
    try:
        # 이미지 저장
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as fp:
            temp_path = fp.name
            Image.fromarray(image).save(temp_path)
        
        # API 호출
        client = Client(VIDEO_API_URL)
        result = client.predict(
            input_image=handle_file(temp_path),
            prompt=prompt if prompt else "Generate natural motion",
            n_prompt="",
            seed=random.randint(0, 9999999),
            use_teacache=True,
            video_length=float(length),
            api_name="/process"
        )
        
        os.unlink(temp_path)
        
        if result and len(result) > 0:
            video_dict = result[0]
            video_path = video_dict.get("video") if isinstance(video_dict, dict) else None
            
            if video_path:
                # Add watermark
                watermarked_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
                if add_watermark_to_video(video_path, watermarked_path):
                    return watermarked_path
                else:
                    return video_path
            
    except Exception as e:
        logging.error(f"Video generation error: {str(e)}")
        return None

# CSS with gradient background
css = """
:root {
    --primary-color: #f8c3cd;
    --secondary-color: #b3e5fc;
    --background-color: #f5f5f7;
    --card-background: #ffffff;
    --text-color: #424242;
    --accent-color: #ffb6c1;
    --success-color: #c8e6c9;
    --warning-color: #fff9c4;
    --shadow-color: rgba(0, 0, 0, 0.1);
    --border-radius: 12px;
}

body, .gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 25%, #f093fb 50%, #4facfe 75%, #00f2fe 100%) !important;
    background-size: 400% 400% !important;
    animation: gradientAnimation 15s ease infinite !important;
}

@keyframes gradientAnimation {
    0% {
        background-position: 0% 50%;
    }
    50% {
        background-position: 100% 50%;
    }
    100% {
        background-position: 0% 50%;
    }
}

.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 20px !important;
}

.panel-box {
    border-radius: var(--border-radius) !important;
    box-shadow: 0 8px 16px var(--shadow-color) !important;
    background-color: var(--card-background) !important;
    padding: 20px !important;
    margin-bottom: 20px !important;
    backdrop-filter: blur(10px) !important;
    background-color: rgba(255, 255, 255, 0.95) !important;
}

#generate-btn, #video-btn {
    background: linear-gradient(135deg, #ff9a9e, #fad0c4) !important;
    font-size: 1.1rem !important;
    padding: 12px 24px !important;
    margin-top: 10px !important;
    width: 100% !important;
    border: none !important;
    color: white !important;
    font-weight: bold !important;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2) !important;
    transition: all 0.3s ease !important;
}

#generate-btn:hover, #video-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(0, 0, 0, 0.25) !important;
}

.main-content {
    min-height: 700px !important;
}

h1 {
    background: linear-gradient(135deg, #667eea, #764ba2) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    text-align: center !important;
    font-size: 2.5rem !important;
    margin-bottom: 30px !important;
    text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.1) !important;
}

.gr-form {
    background: rgba(255, 255, 255, 0.1) !important;
    border-radius: 8px !important;
}

.gr-box {
    border-radius: 8px !important;
}

.gr-padded {
    padding: 15px !important;
}
"""

# Gradio Interface
demo = gr.Blocks(css=css, title="CINE VID Generation")

with demo:
    gr.Markdown("# 🎨 FLUX VIDEO Generation")
    
    # Single tab layout without tabs
    with gr.Row(equal_height=True, elem_classes="main-content"):
        # Input column
        with gr.Column(scale=1):
            with gr.Group(elem_classes="panel-box"):
                gr.Markdown("### 📝 Image Generation Settings")
                
                prompt = gr.Textbox(
                    label="Prompt (Korean/English supported)",
                    placeholder="Describe the image you want to generate...",
                    lines=3
                )
                
                size_preset = gr.Dropdown(
                    choices=list(IMAGE_PRESETS.keys()),
                    value="1:1 Square",
                    label="Size Preset"
                )
                
                with gr.Row():
                    width = gr.Slider(256, 2048, 1024, step=64, label="Width")
                    height = gr.Slider(256, 2048, 1024, step=64, label="Height")
                
                with gr.Row():
                    guidance = gr.Slider(1.0, 20.0, 3.5, step=0.1, label="Guidance Scale")
                    steps = gr.Slider(1, 50, 30, step=1, label="Steps")
                
                seed = gr.Number(label="Seed (-1=random)", value=-1)
                
                generate_btn = gr.Button("🎨 Generate Image", variant="primary", elem_id="generate-btn")
                
            with gr.Group(elem_classes="panel-box"):
                gr.Markdown("### 🎬 Video Generation Settings")
                
                video_prompt = gr.Textbox(
                    label="(Optional) Video Prompt (Enter in English)",
                    placeholder="Describe the motion for the video... (Leave empty for default motion)",
                    lines=2
                )
                
                video_length = gr.Slider(
                    minimum=1, 
                    maximum=4, 
                    value=2, 
                    step=0.5, 
                    label="Video Length (seconds)",
                    info="Choose between 1 to 4 seconds(Full version supports up to 60 seconds)"
                )
                
                video_btn = gr.Button("🎬 Convert to Video", variant="secondary", elem_id="video-btn")
        
        # Output column
        with gr.Column(scale=1):
            with gr.Group(elem_classes="panel-box"):
                gr.Markdown("### 🖼️ Generation Results")
                
                output_image = gr.Image(label="Generated Image", type="numpy")
                output_seed = gr.Textbox(label="Seed Information")
                output_video = gr.Video(label="Generated Video")
    
    # Event connections
    size_preset.change(update_dimensions, [size_preset], [width, height])
    
    generate_btn.click(
        generate_text_to_image,
        [prompt, width, height, guidance, steps, seed],
        [output_image, output_seed]
    )
    
    video_btn.click(
        lambda img, v_prompt, length: generate_video_from_image(img, v_prompt, length) if img is not None else None,
        [output_image, video_prompt, video_length],
        [output_video]
    )

demo.launch()