app.py

import gradio as gr
import os
from lumaai import AsyncLumaAI
import asyncio
import aiohttp
from transformers import pipeline

# 번역 파이프라인 초기화
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

async def get_camera_motions():
    api_key = os.getenv("LMGEN_KEY")
    if not api_key:
        raise gr.Error("LMGEN_KEY environment variable is not set.")
    
    client = AsyncLumaAI(auth_token=api_key)
    try:
        motions = await client.generations.camera_motion.list()
        return [str(motion) for motion in motions]  # Convert each motion to a string
    except Exception as e:
        print(f"Error fetching camera motions: {str(e)}")
        return []

async def generate_video(client, prompt, loop=False, aspect_ratio="16:9", camera_motion=None, extend_id=None, reverse_extend_id=None, interpolate_ids=None, progress=gr.Progress()):
    # aspect_ratio에서 실제 비율만 추출
    aspect_ratio = aspect_ratio.split()[0]

    generation_params = {
        "prompt": prompt,
        "loop": loop,
        "aspect_ratio": aspect_ratio
    }

    if camera_motion:
        generation_params["prompt"] += f" {camera_motion}"

    if extend_id:
        generation_params["keyframes"] = {
            "frame0": {
                "type": "generation",
                "id": extend_id
            }
        }
    elif reverse_extend_id:
        generation_params["keyframes"] = {
            "frame1": {
                "type": "generation",
                "id": reverse_extend_id
            }
        }
    elif interpolate_ids:
        generation_params["keyframes"] = {
            "frame0": {
                "type": "generation",
                "id": interpolate_ids[0]
            },
            "frame1": {
                "type": "generation",
                "id": interpolate_ids[1]
            }
        }

    progress(0, desc="비디오 생성 시작 중...")
    generation = await client.generations.create(**generation_params)
    
    progress(0.1, desc="비디오 생성 시작됨. 완료 대기 중...")
    
    # Poll for completion
    start_time = asyncio.get_event_loop().time()
    while True:
        status = await client.generations.get(id=generation.id)
        if status.state == "completed":
            break
        elif status.state == "failed":
            raise Exception("비디오 생성 실패")
        
        # Update progress based on time elapsed (assuming 60 seconds total)
        elapsed_time = asyncio.get_event_loop().time() - start_time
        progress_value = min(0.1 + (elapsed_time / 60) * 0.8, 0.9)
        progress(progress_value, desc="비디오 생성 중...")
        
        await asyncio.sleep(5)

    progress(0.9, desc="생성된 비디오 다운로드 중...")
    
    # Download the video
    video_url = status.assets.video
    async with aiohttp.ClientSession() as session:
        async with session.get(video_url) as resp:
            if resp.status == 200:
                file_name = f"luma_ai_generated_{generation.id}.mp4"
                with open(file_name, 'wb') as fd:
                    while True:
                        chunk = await resp.content.read(1024)
                        if not chunk:
                            break
                        fd.write(chunk)
    
    progress(1.0, desc="비디오 생성 완료!")
    return file_name, generation.id

async def translate_prompt(prompt):
    try:
        translated = translator(prompt, max_length=512)[0]['translation_text']
        return translated
    except Exception as e:
        print(f"번역 중 오류 발생: {str(e)}")
        return prompt  # 번역 실패 시 원본 프롬프트 반환

async def text_to_video(prompt, loop, aspect_ratio, camera_motion, extend_id, reverse_extend_id, interpolate_id1, interpolate_id2, progress=gr.Progress()):
    api_key = os.getenv("LMGEN_KEY")
    if not api_key:
        raise gr.Error("LMGEN_KEY 환경 변수가 설정되지 않았습니다.")
    
    client = AsyncLumaAI(auth_token=api_key)
    
    try:
        # 프롬프트 번역
        translated_prompt = await translate_prompt(prompt)
        print(f"원본 프롬프트: {prompt}")
        print(f"번역된 프롬프트: {translated_prompt}")

        # aspect_ratio에서 실제 비율만 추출
        aspect_ratio = aspect_ratio.split()[0]

        interpolate_ids = None
        if interpolate_id1 and interpolate_id2:
            interpolate_ids = [interpolate_id1, interpolate_id2]

        video_path, video_id = await generate_video(
            client, translated_prompt, loop, aspect_ratio, camera_motion, 
            extend_id, reverse_extend_id, interpolate_ids, progress
        )
        return video_path, video_id, ""
    except Exception as e:
        return None, None, f"오류 발생: {str(e)}"

async def image_to_video(prompt, image_url, loop, aspect_ratio, camera_motion, progress=gr.Progress()):
    api_key = os.getenv("LMGEN_KEY")
    if not api_key:
        raise gr.Error("LMGEN_KEY 환경 변수가 설정되지 않았습니다.")
    
    client = AsyncLumaAI(auth_token=api_key)
    
    try:
        # 프롬프트 번역
        translated_prompt = await translate_prompt(prompt)
        print(f"원본 프롬프트: {prompt}")
        print(f"번역된 프롬프트: {translated_prompt}")

        # aspect_ratio에서 실제 비율만 추출
        aspect_ratio = aspect_ratio.split()[0]

        generation_params = {
            "prompt": translated_prompt + (f" {camera_motion}" if camera_motion else ""),
            "loop": loop,
            "aspect_ratio": aspect_ratio,
            "keyframes": {
                "frame0": {
                    "type": "image",
                    "url": image_url
                }
            }
        }
        
        progress(0, desc="이미지로부터 비디오 생성 시작 중...")
        generation = await client.generations.create(**generation_params)
        
        progress(0.1, desc="비디오 생성 시작됨. 완료 대기 중...")
        
        # Poll for completion
        start_time = asyncio.get_event_loop().time()
        while True:
            status = await client.generations.get(id=generation.id)
            if status.state == "completed":
                break
            elif status.state == "failed":
                raise Exception("비디오 생성 실패")
            
            # Update progress based on time elapsed (assuming 60 seconds total)
            elapsed_time = asyncio.get_event_loop().time() - start_time
            progress_value = min(0.1 + (elapsed_time / 60) * 0.8, 0.9)
            progress(progress_value, desc="비디오 생성 중...")
            
            await asyncio.sleep(5)

        progress(0.9, desc="생성된 비디오 다운로드 중...")
        
        # Download the video
        video_url = status.assets.video
        async with aiohttp.ClientSession() as session:
            async with session.get(video_url) as resp:
                if resp.status == 200:
                    file_name = f"luma_ai_generated_{generation.id}.mp4"
                    with open(file_name, 'wb') as fd:
                        while True:
                            chunk = await resp.content.read(1024)
                            if not chunk:
                                break
                            fd.write(chunk)
        
        progress(1.0, desc="비디오 생성 완료!")
        return file_name, generation.id, ""
    except Exception as e:
        return None, None, f"오류 발생: {str(e)}"

css = """
footer {
    visibility: hidden;
}
"""

with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
    
    with gr.Tab("텍스트로 비디오 만들기"):
        prompt = gr.Textbox(label="프롬프트 (비디오에 대한 설명을 입력하세요)")
        generate_btn = gr.Button("비디오 생성")
        video_output = gr.Video(label="생성된 비디오")
        video_id_output = gr.Textbox(label="생성된 비디오 ID", visible=True)
        error_output = gr.Textbox(label="오류 메시지", visible=True)
        
        with gr.Accordion("고급 옵션", open=False):
            loop = gr.Checkbox(label="루프 (비디오를 반복 재생할지 설정)", value=False)
            aspect_ratio = gr.Dropdown(
                label="화면 비율", 
                choices=["16:9 (와이드스크린)", "1:1 (정사각형)", "9:16 (세로 영상)", "4:3 (표준)", "3:4 (세로 표준)", "21:9 (울트라와이드)", "9:21 (세로 울트라와이드)"], 
                value="16:9 (와이드스크린)"
            )
            camera_motion = gr.Dropdown(label="카메라 모션 (카메라 움직임 효과 선택)")
            extend_id = gr.Textbox(label="확장할 비디오 ID (기존 비디오를 이어서 생성할 때 입력)")
            reverse_extend_id = gr.Textbox(label="역방향 확장할 비디오 ID (기존 비디오의 앞부분을 생성할 때 입력)")
            with gr.Row():
                interpolate_id1 = gr.Textbox(label="보간 비디오 ID 1 (두 비디오 사이를 보간할 때 첫 번째 비디오 ID)")
                interpolate_id2 = gr.Textbox(label="보간 비디오 ID 2 (두 비디오 사이를 보간할 때 두 번째 비디오 ID)")
        
        generate_btn.click(
            text_to_video,
            inputs=[prompt, loop, aspect_ratio, camera_motion, extend_id, reverse_extend_id, interpolate_id1, interpolate_id2],
            outputs=[video_output, video_id_output, error_output]
        )
    
    with gr.Tab("이미지로 비디오 만들기"):
        img_prompt = gr.Textbox(label="프롬프트 (이미지를 바탕으로 생성할 비디오에 대한 설명)")
        img_url = gr.Textbox(label="이미지 URL (변환하고자 하는 이미지의 웹 주소)")
        img_generate_btn = gr.Button("이미지로 비디오 생성")
        img_video_output = gr.Video(label="생성된 비디오")
        img_video_id_output = gr.Textbox(label="생성된 비디오 ID", visible=True)
        img_error_output = gr.Textbox(label="오류 메시지", visible=True)
        
        with gr.Accordion("고급 옵션", open=False):
            img_loop = gr.Checkbox(label="루프 (비디오를 반복 재생할지 설정)", value=False)
            img_aspect_ratio = gr.Dropdown(
                label="화면 비율", 
                choices=["16:9 (와이드스크린)", "1:1 (정사각형)", "9:16 (세로 영상)", "4:3 (표준)", "3:4 (세로 표준)", "21:9 (울트라와이드)", "9:21 (세로 울트라와이드)"], 
                value="16:9 (와이드스크린)"
            )
            img_camera_motion = gr.Dropdown(label="카메라 모션 (카메라 움직임 효과 선택)")
        
        img_generate_btn.click(
            image_to_video,
            inputs=[img_prompt, img_url, img_loop, img_aspect_ratio, img_camera_motion],
            outputs=[img_video_output, img_video_id_output, img_error_output]
        )

    async def update_camera_motions():
        try:
            motions = await get_camera_motions()
            return gr.update(choices=motions), gr.update(choices=motions)
        except Exception as e:
            print(f"카메라 모션 업데이트 중 오류 발생: {str(e)}")
            return gr.update(choices=[]), gr.update(choices=[])

    demo.load(update_camera_motions, outputs=[camera_motion, img_camera_motion])

demo.queue().launch(auth=("gini","pick"))