animily_music/generate_image.py

"""Animily Music - FLUX 이미지 생성 (ComfyUI API)

신카이 마코토 스타일 이미지 1장 생성 (1920x1080, 16:9 가로).
"""
import glob
import json
import os
import shutil
import signal
import subprocess
import time

import requests

from config import COMFYUI_URL, COMFYUI_DIR, OUTPUT_DIR, PROMPTS_DIR


def _start_comfyui():
    """ComfyUI FLUX 서버 시작 (8189 포트)"""
    try:
        requests.get(COMFYUI_URL, timeout=3)
        print("  [FLUX] ComfyUI 이미 실행 중", flush=True)
        return
    except Exception:
        pass

    # 기존 8188 (WAN) 종료
    try:
        requests.get("http://localhost:8188", timeout=2)
        result = subprocess.run(["fuser", "8188/tcp"], capture_output=True, text=True)
        for p in result.stdout.strip().split():
            if p.strip().isdigit():
                os.kill(int(p.strip()), signal.SIGTERM)
        time.sleep(5)
    except Exception:
        pass

    # FLUX ComfyUI 시작
    print("  [FLUX] ComfyUI 시작 중...", flush=True)
    subprocess.Popen(
        [f"{COMFYUI_DIR}/venv/bin/python3", "main.py",
         "--listen", "0.0.0.0", "--port", "8189", "--bf16-vae", "--lowvram"],
        cwd=COMFYUI_DIR,
        stdout=open("/tmp/comfyui_flux_music.log", "a"),
        stderr=open("/tmp/comfyui_flux_music.log", "a"),
    )
    for _ in range(120):
        try:
            requests.get(COMFYUI_URL, timeout=2)
            print("  [FLUX] ComfyUI 시작 완료", flush=True)
            return
        except Exception:
            time.sleep(1)
    raise RuntimeError("ComfyUI 시작 실패 (120초 타임아웃)")


def _stop_comfyui():
    """ComfyUI 종료 및 GPU 메모리 해제"""
    try:
        result = subprocess.run(["fuser", "8189/tcp"], capture_output=True, text=True)
        for p in result.stdout.strip().split():
            if p.strip().isdigit():
                os.kill(int(p.strip()), signal.SIGTERM)
        time.sleep(3)
        # 강제 종료
        result2 = subprocess.run(["fuser", "8189/tcp"], capture_output=True, text=True)
        for p in result2.stdout.strip().split():
            if p.strip().isdigit():
                try:
                    os.kill(int(p.strip()), signal.SIGKILL)
                except Exception:
                    pass
        time.sleep(5)
        print("  [FLUX] ComfyUI 종료 + GPU 메모리 해제 완료", flush=True)
    except Exception as e:
        print(f"  [FLUX] ComfyUI 종료 중 오류: {e}", flush=True)


def generate_image(prompt, output_path):
    """FLUX로 이미지 생성 (1920x1080)

    Args:
        prompt: 영어 이미지 프롬프트
        output_path: 출력 파일 경로

    Returns:
        bool: 성공 여부
    """
    _start_comfyui()

    # 이전 출력 정리
    for f in glob.glob(f"{COMFYUI_DIR}/output/flux_gen_*.png"):
        os.remove(f)

    # FLUX.2 워크플로우 (1920x1080)
    workflow = {
        "1": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": "flux2-dev-Q8_0.gguf"}},
        "2": {"class_type": "LoraLoader", "inputs": {
            "model": ["1", 0], "clip": ["3", 0],
            "lora_name": "flux2_turbo_comfy.safetensors",
            "strength_model": 1.0, "strength_clip": 0.0,
        }},
        "3": {"class_type": "CLIPLoader", "inputs": {
            "clip_name": "mistral_3_small_flux2_bf16.safetensors", "type": "flux2",
        }},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": prompt}},
        "5": {"class_type": "EmptyLatentImage", "inputs": {
            "width": 1024, "height": 576, "batch_size": 1,
        }},
        "6": {"class_type": "KSampler", "inputs": {
            "model": ["2", 0], "positive": ["4", 0], "negative": ["4", 0],
            "latent_image": ["5", 0],
            "seed": int(time.time()) % 999999, "steps": 8, "cfg": 1.0,
            "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0,
        }},
        "7": {"class_type": "VAELoader", "inputs": {"vae_name": "flux2-ae.safetensors"}},
        "8": {"class_type": "VAEDecodeTiled", "inputs": {
            "vae": ["7", 0], "samples": ["6", 0],
            "tile_size": 512, "overlap": 64, "temporal_size": 64, "temporal_overlap": 8,
        }},
        "9": {"class_type": "SaveImage", "inputs": {
            "images": ["8", 0], "filename_prefix": "flux_gen",
        }},
    }

    print(f"  [FLUX] 이미지 생성 중...", flush=True)
    t0 = time.time()

    try:
        resp = requests.post(f"{COMFYUI_URL}/prompt", json={"prompt": workflow}, timeout=10)
        if resp.status_code != 200:
            print(f"  [FLUX] 워크플로우 에러: {resp.text[:200]}", flush=True)
            return False

        pid = resp.json().get("prompt_id")

        # 완료 대기 (최대 15분)
        while (time.time() - t0) < 900:
            h = requests.get(f"{COMFYUI_URL}/history/{pid}", timeout=5).json()
            if pid in h:
                status = h[pid].get("status", {}).get("status_str", "")
                if status == "success":
                    pngs = sorted(glob.glob(f"{COMFYUI_DIR}/output/flux_gen_*.png"))
                    if pngs:
                        shutil.copy2(pngs[-1], output_path)
                        for p in pngs:
                            os.remove(p)
                        elapsed = time.time() - t0
                        print(f"  [FLUX] 이미지 생성 완료: {elapsed:.0f}초", flush=True)
                        return True
                    return False
                elif status == "error":
                    print(f"  [FLUX] 생성 에러", flush=True)
                    return False
            time.sleep(5)

        print(f"  [FLUX] 타임아웃 (15분)", flush=True)
        return False

    except Exception as e:
        print(f"  [FLUX] 오류: {e}", flush=True)
        return False


def generate_thumbnail(image_path, thumbnail_path):
    """1280x720 썸네일 생성"""
    cmd = [
        "ffmpeg", "-y", "-i", image_path,
        "-vf", "scale=1280:720",
        "-q:v", "2", thumbnail_path,
    ]
    subprocess.run(cmd, capture_output=True)
    return os.path.exists(thumbnail_path)


def load_image_prompt(animal_type="dog"):
    """이미지 프롬프트 로드"""
    path = os.path.join(PROMPTS_DIR, "image_prompts.json")
    with open(path) as f:
        data = json.load(f)
    prompts = data.get(animal_type, data.get("dog"))
    return prompts[int(time.time()) % len(prompts)]


if __name__ == "__main__":
    import sys
    animal = sys.argv[1] if len(sys.argv) > 1 else "dog"
    prompt = load_image_prompt(animal)
    out = os.path.join(OUTPUT_DIR, f"test_image_{animal}.png")
    ok = generate_image(prompt, out)
    _stop_comfyui()
    print(f"결과: {ok}, {out}")