import os
import gc
import torch
import gradio as gr
from PIL import Image
from diffusers import DiffusionPipeline

# --- [초핵심 크리티컬 세팅] ---
# PyTorch가 safetensors를 로드할 때 스토리지 파일 매핑(mmap)으로 RAM을 중복 할당하는 것을 방지
os.environ["SAFETENSORS_FAST_GPU"] = "0" 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🖥️ [System] 현재 감지된 실행 디바이스: {DEVICE.upper()}")

def clean_memory():
    """가비지 컬렉션을 아주 강력하게 수동 트리거"""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

# --- 2. 파이프라인 초기화 로더 (RAM 96GB 세이브 모드) ---
print("🚀 [System] 스튜디오 파이프라인 최적화 초기화 시작...")

# [Step 1] FireRed Image Edit 로드
print("📦 [1/2] FireRed Image Edit 모델 로드 중...")
pipe_edit = DiffusionPipeline.from_pretrained(
    "FireRedTeam/FireRed-Image-Edit-1.1", 
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
    low_cpu_mem_usage=True, # RAM 버스트 방어
)

if DEVICE == "cuda":
    pipe_edit.to("cuda")
else:
    # CPU 환경일 때 RAM 용량이 터지지 않도록 레이어 단위 CPU 오프로드 강제
    pipe_edit.enable_model_cpu_offload()

print("✅ [System] FireRed Image Edit 로드 완료.")
clean_memory() # 1번 모델 로드 후 찌꺼기 즉시 삭제

# [Step 2] Wan 2.2 14B 로드
print("📦 [2/2] Wan 2.2 14B 모델 로드 중...")
# CPU 환경일 때는 fp8 커널이 에러를 내므로 전용 옵션을 가변적으로 매핑
wan_kwargs = {
    "torch_dtype": torch.bfloat16 if DEVICE == "cuda" else torch.float32,
    "low_cpu_mem_usage": True,
    "device_map": "auto"
}
if DEVICE == "cuda":
    wan_kwargs["variant"] = "fp8"

pipe_wan = DiffusionPipeline.from_pretrained(
    "Wan-Video/Wan2.2-I2V-14B", 
    **wan_kwargs
)

if DEVICE == "cuda":
    pipe_wan.enable_model_cpu_offload()
else:
    # CPU 빌드일 때 96GB RAM 세이브를 위한 최종 병기
    # 파이프라인 내부의 중복 컴포넌트를 청소
    pipe_wan.to("cpu")

print("✅ [System] Wan 2.2 14B 로드 완료.")
clean_memory()


# --- 3. 비즈니스 로직 ---
def process_studio_pipeline(input_image, edit_prompt, video_prompt, num_frames=81, progress=gr.Progress()):
    if input_image is None:
        raise gr.Error("먼저 원본 이미지를 업로드해주세요!")
    
    # [Step 1] FireRed 이미지 편집
    progress(0.1, desc="FireRed로 이미지 편집 중...")
    pil_img = Image.fromarray(input_image).convert("RGB")
    
    edited_image = pipe_edit(
        prompt=edit_prompt,
        image=pil_img,
        num_inference_steps=20 if DEVICE == "cuda" else 3 # CPU일 땐 최소 연산
    ).images[0]
    
    clean_memory()
    
    # [Step 2] Wan 2.2 14B 비디오 생성
    progress(0.5, desc="Wan 2.2로 비디오 생성 중...")
    
    video_frames = pipe_wan(
        prompt=video_prompt,
        image=edited_image,
        num_frames=int(num_frames) if DEVICE == "cuda" else 16, # CPU 연산 최소화
        num_inference_steps=30 if DEVICE == "cuda" else 3,
        guidance_scale=6.0
    ).frames[0]
    
    clean_memory()
    
    # [Step 3] 비디오 파일 저장
    progress(0.9, desc="비디오 렌더링 중...")
    output_video_path = "studio_output.mp4"
    
    from diffusers.utils import export_to_video
    export_to_video(video_frames, output_video_path, fps=16)
    
    return edited_image, output_video_path


# --- 4. Gradio UI Layout ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎬 iE2V Fast Studio (RAM Dynamic Build)")
    if DEVICE == "cpu":
        gr.Warning("⚠️ 인프라가 CPU 모드로 제한되어 RAM 96GB 크래시 방어 모드가 작동 중입니다. 정상 추론을 위해서는 Space Settings에서 GPU 할당을 꼭 확보해 주세요.")
        
    with gr.Row():
        with gr.Column():
            input_img_slot = gr.Image(label="1. 원본 이미지 업로드", type="numpy")
            edit_prompt_slot = gr.Textbox(label="2. 이미지 편집 프롬프트 (FireRed)", placeholder="배경을 밤거리로 변경...")
            video_prompt_slot = gr.Textbox(label="3. 비디오 연출 프롬프트 (Wan 2.2)", placeholder="카메라가 서서히 줌인되며...")
            
            with gr.Accordion("🎥 고급 세팅", open=False):
                frames_slider = gr.Slider(minimum=16, maximum=81, step=4, value=81 if DEVICE == "cuda" else 16, label="프레임 수")
                
            submit_btn = gr.Button("⚡ 융합 생성 시작", variant="primary")
            
        with gr.Column():
            edited_img_output = gr.Image(label="편집된 이미지 결과")
            video_output = gr.Video(label="최종 생성 비디오")
            
    submit_btn.click(
        fn=process_studio_pipeline,
        inputs=[input_img_slot, edit_prompt_slot, video_prompt_slot, frames_slider],
        outputs=[edited_img_output, video_output]
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)