run_server.py

import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))

def add_directory_and_subdirectories_to_sys_path(base_path):
    """
    지정된 디렉터리와 하위 디렉터리를 sys.path에 추가합니다.
    :param base_path: 기준 디렉터리의 절대 경로
    """
    for root, dirs, files in os.walk(base_path):
        if root not in sys.path:  # 중복 방지
            sys.path.append(root)

# 절대 경로로 기본 디렉터리 지정
base_directory = "/root/app"
# 경로 및 하위 디렉터리 모두 sys.path에 추가
add_directory_and_subdirectories_to_sys_path(base_directory)

import base64
from flask import Flask, request, jsonify
from PIL import Image, ImageDraw, ImageFont
import io
import torch
import json
import numpy as np
import random
import comfy.model_management
import comfy.sd
import comfy.utils
import folder_paths
import node_helpers
from nodes import CLIPTextEncode, VAEDecode, CheckpointLoaderSimple, LoraLoader, common_ksampler, LoadImage
from datetime import datetime
import threading

from custom_nodes.ComfyUI_segment_anything2.nodes import DownloadAndLoadSAM2Model, Sam2Segmentation, \
    Florence2toCoordinates
from custom_nodes.ComfyUI_Florence2.nodes import Florence2ModelLoader, DownloadAndLoadFlorence2Model, Florence2Run
import time

from custom_nodes.Fooocus_Nodes.py.modules import model_loader
from custom_nodes.Fooocus_Nodes.py.fooocusNodes import (
    FooocusLoader,
    FooocusPreKSampler,
    FooocusControlnet,
    FooocusKsampler
)
from custom_nodes.keyvisual.wcoh import NODE_CLASS_MAPPINGS as keyvisual_nodes

from custom_nodes.ComfyUI_essentials.image import ImageResize
from custom_nodes.efficiency_nodes_comfyui.efficiency_nodes import TSC_ImageOverlay
import torch.nn.functional as F

import cv2
import insightface
from gfpgan.utils import GFPGANer
import tempfile

app = Flask(__name__)

# 샘플러와 스케줄러 옵션 정의
SAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
                 "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral",
                 "dpmpp_sde", "dpmpp_2m", "ddim", "uni_pc", "uni_pc_bh2"]

SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]

# Welcome Poster classes
class FaceChanger:
    def __init__(self):
        self.face_swapper_model_path = '/data/public/models/inswapper_128.onnx'
        self.face_enhancer_model_path = '/data/public/models/GFPGANv1.4.pth'
        self.face_analyzer = insightface.app.FaceAnalysis(name='buffalo_l', providers=["CUDAExecutionProvider"])
        self.face_analyzer.prepare(ctx_id=0)
        self.face_swapper = insightface.model_zoo.get_model(self.face_swapper_model_path, providers=["CUDAExecutionProvider"])
        self.face_enhancer = GFPGANer(model_path=self.face_enhancer_model_path, upscale=1, device="cpu")

    def get_face(self, image_path):
        src = cv2.imread(image_path)
        faces = self.face_analyzer.get(src)
        return faces[0] if faces else None

    def swap_face(self, source_face, target_face, target_frame):
        return self.face_swapper.get(target_frame, target_face, source_face, paste_back=True)

    def enhance_face(self, image):
        """얼굴을 선명하게 하는 단계 추가"""
        _, _, enhanced_image = self.face_enhancer.enhance(image, has_aligned=False, only_center_face=False, paste_back=True)
        return enhanced_image

    def run(self, source_image_path, target_image_path):
        source_face = self.get_face(source_image_path)
        target_frame = cv2.imread(target_image_path)
        target_face = self.get_face(target_image_path)
        if not source_face or not target_face:
            raise ValueError("[ERROR] 얼굴을 감지할 수 없습니다.")
        swapped_face = self.swap_face(source_face, target_face, target_frame)

        enhanced_face = self.enhance_face(swapped_face)
        return enhanced_face
        

class WCOHRandomImage:
    def __init__(self, folder_path):
        self.folder_path = folder_path

    def select_random_image(self):
        valid_extensions = (".png", ".jpg", ".jpeg", ".bmp", ".gif")
        image_files = [f for f in os.listdir(self.folder_path) if f.lower().endswith(valid_extensions)]
        if not image_files:
            raise ValueError(f"No valid image files in {self.folder_path}")
        return os.path.join(self.folder_path, random.choice(image_files))

class WCOHTextOnImage:
    def __init__(self):
        self.team_font_path = "/root/app/custom_nodes/wcoh/Jalnan2TTF.ttf"
        self.message_font_path = "/root/app/custom_nodes/wcoh/Incheon.ttf"

    def add_text(self, image, team_name, user_name, message, concept="1"):
        draw = ImageDraw.Draw(image)
        if concept == "3":
            x_offset, y_offset, xy_offset, font_scale = (120, 740, 717, -20)
        elif concept == "1":
            x_offset, y_offset, xy_offset, font_scale = (0, 720, 720, 10)
        else:
            x_offset, y_offset, xy_offset, font_scale = (0, 0, 0, 0)

        team_font = ImageFont.truetype(self.team_font_path, 80 + font_scale)
        message_font = ImageFont.truetype(self.message_font_path, 45)

        draw.text((52 + x_offset, 122 + y_offset), team_name, font=team_font, fill="black")
        draw.text((50 + x_offset, 120 + y_offset), team_name, font=team_font, fill="white")
        draw.text((72 + x_offset, 222 + xy_offset), user_name, font=team_font, fill="black")
        draw.text((70 + x_offset, 220 + xy_offset), user_name, font=team_font, fill="white")

        if concept != "3":
            image_width, _ = image.size
            bbox = message_font.getbbox(message)
            text_width = bbox[2] - bbox[0]
            message_x = (image_width - text_width) // 2
            draw.text((message_x + 1, 1101), message, font=message_font, fill="black")
            draw.text((message_x, 1100), message, font=message_font, fill="white")


        return image


# 모델 상태를 관리하기 위한 클래스
class ModelState:
    def __init__(self):
        self.lock = threading.Lock()
        self.is_generating = False

# 전역 상태 관리
model_states = {
    "A": ModelState(),
    "B": ModelState(),
    "SEGMENT_A": ModelState(),
    "SEGMENT_B": ModelState(),
    "KEYVISUAL": ModelState(),
    "WELCOME_POSTER":ModelState()
}

# 전역 변수로 모델들을 초기화
model = None
clip = None
vae = None
model_b = None
clip_b = None
vae_b = None
sam2_model_a = None
F_model_a = None
sam2_model_b = None
F_model_b = None


def load_image_base64_to_numpy(base64_image):
    image_data = base64.b64decode(base64_image)
    image = Image.open(io.BytesIO(image_data)).convert("RGB")
    image_np = np.array(image)
    image_np = image_np[None, ...]  # (H, W, C) -> (B, H, W, C) 변환
    return image_np


def initialize_models(model_type="A"):

    global model, clip, vae, model_b, clip_b, vae_b
    global sam2_model_a, F_model_a, sam2_model_b, F_model_b
    global global_face_changer  # Welcome Poster 모델을 위한 전역 변수 추가

    if model_type == "A":
        try:
            with torch.no_grad():
                checkpoint_loader = CheckpointLoaderSimple()
                model, clip, vae = checkpoint_loader.load_checkpoint(ckpt_name="flux1-schnell-fp8.safetensors")

                lora_loader = LoraLoader()
                model, clip = lora_loader.load_lora(
                    model=model,
                    clip=clip,
                    lora_name="FluxDFaeTasticDetails.safetensors",
                    strength_model=1.0,
                    strength_clip=1.0
                )
                print(f"Model Loaded Completed with model_type '{model_type}'")

        except Exception as e:
            print(f"Error initializing models: {str(e)}")
            raise

    elif model_type == "B":
        try:
            with torch.no_grad():
                checkpoint_loader = CheckpointLoaderSimple()
                model_b, clip_b, vae_b = checkpoint_loader.load_checkpoint(ckpt_name="flux1-schnell-fp8.safetensors")

                lora_loader = LoraLoader()
                model_b, clip_b = lora_loader.load_lora(
                    model=model_b,
                    clip=clip_b,
                    lora_name="FluxDFaeTasticDetails.safetensors",
                    strength_model=1.0,
                    strength_clip=1.0
                )
                print(f"Model Loaded Completed with model_type '{model_type}'")

        except Exception as e:
            print(f"Error initializing models: {str(e)}")
            raise



    elif model_type == "SEGMENT_A":

        try:

            # SAM2 모델 A 로드
            sam2_loader = DownloadAndLoadSAM2Model()
            global sam2_model_a
            sam2_model_a = sam2_loader.loadmodel(
                model="sam2_hiera_base_plus.safetensors",
                segmentor="single_image",
                device="cuda",
                precision="bf16"
            )

            # Florence2 모델 A 로드

            global F_model_a
            F_model_a = DownloadAndLoadFlorence2Model().loadmodel(
                model='microsoft/Florence-2-base',
                precision="fp16",
                attention="sdpa"
            )

            print(f"Segmentation models A loaded successfully")


        except Exception as e:
            print(f"Error initializing segmentation models A: {str(e)}")

            raise


    elif model_type == "SEGMENT_B":

        try:

            # SAM2 모델 B 로드

            sam2_loader = DownloadAndLoadSAM2Model()
            global sam2_model_b
            sam2_model_b = sam2_loader.loadmodel(
                model="sam2_hiera_base_plus.safetensors",
                segmentor="single_image",
                device="cuda",
                precision="bf16"
            )

            # Florence2 모델 B 로드
            global F_model_b
            F_model_b = DownloadAndLoadFlorence2Model().loadmodel(
                model='microsoft/Florence-2-base',
                precision="fp16",
                attention="sdpa"
            )

            print(f"Segmentation models B loaded successfully")


        except Exception as e:
            print(f"Error initializing segmentation models B: {str(e)}")
            raise

    elif model_type == "KEYVISUAL":
        try:
            # 커스텀 모델 경로 추가
            folder_paths.add_model_folder_path('checkpoints', '/data/public/models/checkpoint')
            folder_paths.add_model_folder_path('vae', '/data/public/models/vae')
            # Fooocus Loader 초기화
            loader = FooocusLoader()
            global keyvisual_pipe
            keyvisual_pipe = loader.fooocus_loader(
                positive="",  # 초기화시에는 빈 프롬프트
                negative="",
                base_model_name="DreamShaperXL_Turbo_v2_1.safetensors",  # 기본 모델
                vae_name="sdxl_vae.safetensors",  # 기본 VAE
                refiner_model_name="None",
                refiner_switch=0.5,
                refiner_swap_method="joint",
                clip_skip=2,
                resolution="832 x 1152",
                image_number=1
            )["result"][0]

            # PreKSampler로 기본 설정 초기화
            preksampler = FooocusPreKSampler()
            result = preksampler.fooocus_preKSampler(
                pipe=keyvisual_pipe,
                steps=8,
                cfg=4,
                sampler_name="dpmpp_sde",
                scheduler="karras",
                denoise=1.0,
                settings="Simple",
                controlnet_softness=0.25,
                sharpness=2.0,
                adaptive_cfg=7.0,
                adm_scaler_positive=1.5,
                adm_scaler_negative=0.8,
                adm_scaler_end=0.3,
                freeu_enabled=False,
                seed=0
            )

            global keyvisual_model
            keyvisual_model = result["result"][1]

            print(f"Keyvisual model loaded successfully")

        except Exception as e:
            print(f"Error initializing keyvisual models: {str(e)}")
            raise

    elif model_type == "WELCOME_POSTER":  # Welcome Poster 초기화 추가
        try:
            global global_face_changer
            global_face_changer = FaceChanger()
            print("Welcome Poster models initialized successfully!")
        except Exception as e:
            print(f"Error initializing Welcome Poster models: {str(e)}")
            raise

def create_empty_latent(width=1024, height=1024, batch_size=1):
    """SD3용 빈 레이턴트 이미지 생성"""
    device = comfy.model_management.intermediate_device()
    latent = torch.zeros([batch_size, 16, height // 8, width // 8], device=device)
    return {"samples": latent}


def sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0):
    """KSampler 실행"""
    with torch.no_grad():
        result = common_ksampler(
            model=model,
            seed=seed,
            steps=steps,
            cfg=cfg,
            sampler_name=sampler_name,
            scheduler=scheduler,
            positive=positive,
            negative=negative,
            latent=latent,
            denoise=denoise,
            disable_noise=False,
            start_step=None,
            last_step=None,
            force_full_denoise=True
        )[0]
        return result["samples"]


def save_json_with_timestamp(data, seed, request_type):
    """JSON 파일을 타임스탬프와 시드값으로 저장"""
    # request_type에 따라 다른 디렉토리 사용
    log_dir = f"/data/public/logs/{request_type.lower()}"
    if not os.path.exists(log_dir):
        os.makedirs(log_dir, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{timestamp}_{seed}_server.json"
    filepath = os.path.join(log_dir, filename)

    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    print(f"Log saved as: {filepath}")

    return filepath


def save_image_with_timestamp(image, seed, request_type):
    """이미지를 타임스탬프와 시드값으로 저장"""
    image_dir = f"/data/public/generated_images/{request_type.lower()}"
    if not os.path.exists(image_dir):
        os.makedirs(image_dir, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{timestamp}_{seed}_server.png"
    filepath = os.path.join(image_dir, filename)

    image.save(filepath)
    print(f"Image saved as: {filepath}")
    return filepath

def save_error_json(error_message, request_type):
    """에러 메시지를 JSON 파일로 저장"""
    log_dir = f"/data/public/logs/{request_type.lower()}/error"
    if not os.path.exists(log_dir):
        os.makedirs(log_dir, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    error_data = {
        'status': 'error',
        'message': error_message,
        'request_type': request_type,
        'timestamp': timestamp
    }
    
    filename = f"{timestamp}_error_server.json"
    filepath = os.path.join(log_dir, filename)

    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(error_data, f, ensure_ascii=False, indent=2)

    print(f"Error log saved as: {filepath}")

    return error_data


def generation_image_by_model(
        model_var_name: str,
        clip_var_name: str,
        vae_var_name: str
):
    model = globals()[model_var_name] if model_var_name in globals() else None
    clip = globals()[clip_var_name] if model_var_name in globals() else None
    vae = globals()[vae_var_name] if model_var_name in globals() else None

    data = request.get_json()
    request_type = data.get('request_type', None)

    try:
        with torch.no_grad():
            prompt = data.get('prompt', '')
            negative_prompt = data.get('negative_prompt', '')
            width = data.get('width', 1024)
            height = data.get('height', 1024)
            guidance_value = data.get('guidance_scale', 3.0)
            steps = data.get('steps', 30)
            cfg = data.get('cfg_scale', 1.0)
            seed = data.get('seed', random.randint(0, 2 ** 32 - 1))
            sampler_name = data.get('sampler_name', 'euler')
            scheduler = data.get('scheduler', 'normal')
            denoise = data.get('denoise', 1.0)

            print(f"Using seed: {seed}")

            # CLIP Text Encode
            clip_encode = CLIPTextEncode()
            positive_cond = clip_encode.encode(clip, prompt)[0]
            negative_cond = clip_encode.encode(clip, negative_prompt)[0]

            # FluxGuidance 적용
            positive_cond = node_helpers.conditioning_set_values(positive_cond, {"guidance": guidance_value})
            # SD3용 빈 레이턴트 생성
            latent = create_empty_latent(width=width, height=height)

            # Sampling 실행
            samples = sample(
                model=model,
                seed=seed,
                steps=steps,
                cfg=cfg,
                sampler_name=sampler_name,
                scheduler=scheduler,
                positive=positive_cond,
                negative=negative_cond,
                latent=latent,
                denoise=denoise
            )

            # VAE Decode
            vae_decoder = VAEDecode()
            decoded = vae_decoder.decode(vae, {"samples": samples})[0]

            # 텐서 형태 변환
            decoded = decoded.cpu().numpy()
            if len(decoded.shape) == 5:
                decoded = decoded[0, 0]
            elif len(decoded.shape) == 4:
                decoded = decoded[0]

            # uint8로 변환
            decoded = np.clip(decoded * 255, 0, 255).astype(np.uint8)

            # 이미지로 변환
            image = Image.fromarray(decoded)

            # 이미지 파일 저장
            image_filepath = save_image_with_timestamp(image, seed, request_type)

            # base64 포맷 변환
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()

            # 응답 데이터 생성
            response_data = {
                'status': 'success',
                'image': img_str,
                'parameters': {
                    'prompt': prompt,
                    'negative_prompt': negative_prompt,
                    'width': width,
                    'height': height,
                    'guidance_scale': guidance_value,
                    'steps': steps,
                    'cfg_scale': cfg,
                    'seed': seed,
                    'sampler_name': sampler_name,
                    'scheduler': scheduler,
                    'denoise': denoise
                },
                'request_type': request_type,
                
            }

            # JSON 파일 저장
            json_filepath = save_json_with_timestamp(response_data, seed, request_type)  # request_type 추가
            
            return response_data

    except Exception as e:
        raise e


# API 엔드포인트들
@app.route('/v1/models/t2i:state', methods=['GET'])
def get_generation_state():
    """t2i 모델 A의 현재 생성 상태를 반환"""
    with model_states['A'].lock:
        state = "1" if model_states['A'].is_generating else "0"
    return jsonify({"status": state})

@app.route('/v1/models/t2i:state2', methods=['GET'])
def get_generation_state2():
    """t2i 모델 B의 현재 생성 상태를 반환"""
    with model_states['B'].lock:
        state = "1" if model_states['B'].is_generating else "0"
    return jsonify({"status": state})

@app.route('/v1/models/segment:state', methods=['GET'])
def get_segment_state():
    """세그멘테이션 모델 A의 현재 상태를 반환"""
    with model_states["SEGMENT_A"].lock:
        state = "1" if model_states["SEGMENT_A"].is_generating else "0"
    return jsonify({"status": state})

@app.route('/v1/models/segment:state2', methods=['GET'])
def get_segment_state2():
    """세그멘테이션 모델 B의 현재 상태를 반환"""
    with model_states["SEGMENT_B"].lock:
        state = "1" if model_states["SEGMENT_B"].is_generating else "0"
    return jsonify({"status": state})

@app.route('/v1/models/keyvisual:state', methods=['GET'])
def get_keyvisual_state():
    """키비주얼 모델의 현재 상태를 반환"""
    with model_states["KEYVISUAL"].lock:
        state = "1" if model_states["KEYVISUAL"].is_generating else "0"
    return jsonify({"status": state})

@app.route('/v1/models/welcomeposter:state', methods=['GET'])
def get_welcome_poster_state():
    """Welcome Poster 모델의 현재 상태를 반환"""
    with model_states["WELCOME_POSTER"].lock:
        state = "1" if model_states["WELCOME_POSTER"].is_generating else "0"
    return jsonify({"status": state})

@app.route('/v1/models/segment:point', methods=['POST'])
def segmentation_point():
    """세그멘테이션 A 모델의 포인트 기반 세그멘테이션"""
    state = model_states["SEGMENT_A"]

    with state.lock:
        if state.is_generating:
            return jsonify({'status': 'error', 'message': 'Segmentation engine A is busy'}), 503
        state.is_generating = True

    try:
        #time.sleep(5)
        data = request.get_json()
        #data['request_type'] = 'A'  # request_type 강제 설정
        request_type = data.get('request_type')
        base64_image = data.get('image')
        image_path = data.get('image_path', 'no image path')

        if not base64_image:
                error_response = save_error_json(str(e), request_type)
                return jsonify(error_response), 400

        image_np = load_image_base64_to_numpy(base64_image)
        image_tensor = torch.from_numpy(image_np).float()

        coordinates_positive = str(data.get('coordinates_positive')) if data.get('coordinates_positive') else None
        coordinates_negative = str(data.get('coordinates_negative')) if data.get('coordinates_negative') else None

        segmenter = Sam2Segmentation()
        mask_output = segmenter.segment(
            image=image_tensor,
            sam2_model=sam2_model_a[0],
            keep_model_loaded=True,
            coordinates_positive=coordinates_positive if coordinates_positive else None,
            coordinates_negative=coordinates_negative if coordinates_negative else None,
        )

        mask_image = mask_output[0].cpu().numpy()
        mask_image = np.squeeze(mask_image)
        mask_image = Image.fromarray((mask_image * 255).astype('uint8'))

        buffered = io.BytesIO()
        mask_image.save(buffered, format="PNG")
        mask_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        randomint = random.randint(1000000000, 9999999999)

        json_data = {
            'status': 'success',
            'image_path': image_path,
            'coordinates_positive': coordinates_positive,
            'coordinates_negative': coordinates_negative,
            'timestamp': timestamp,
            'random_int': randomint,
            'mask': mask_base64,
            'request_type': request_type
        }

        json_filepath = save_json_with_timestamp(json_data, randomint, request_type)


        return jsonify(json_data)

    except Exception as e:
        error_response = save_error_json(str(e), request_type)
        return jsonify(error_response), 500
    finally:
        with state.lock:
            state.is_generating = False

@app.route('/v1/models/segment:point2', methods=['POST'])
def segmentation_point2():
    """세그멘테이션 B 모델의 포인트 기반 세그멘테이션"""
    state = model_states["SEGMENT_B"]

    with state.lock:
        if state.is_generating:
            return jsonify({'status': 'error', 'message': 'Segmentation engine B is busy'}), 503
        state.is_generating = True

    try:
        #time.sleep(5)
        data = request.get_json()
        #data['request_type'] = 'B'  # request_type 강제 설정
        request_type = data.get('request_type')
        base64_image = data.get('image')
        image_path = data.get('image_path', 'no image path')

        if not base64_image:
            error_response = save_error_json(str(e), request_type)
            return jsonify(error_response), 400

        image_np = load_image_base64_to_numpy(base64_image)
        image_tensor = torch.from_numpy(image_np).float()

        coordinates_positive = str(data.get('coordinates_positive')) if data.get('coordinates_positive') else None
        coordinates_negative = str(data.get('coordinates_negative')) if data.get('coordinates_negative') else None

        segmenter = Sam2Segmentation()
        mask_output = segmenter.segment(
            image=image_tensor,
            sam2_model=sam2_model_b[0],
            keep_model_loaded=True,
            coordinates_positive=coordinates_positive if coordinates_positive else None,
            coordinates_negative=coordinates_negative if coordinates_negative else None,
        )

        mask_image = mask_output[0].cpu().numpy()
        mask_image = np.squeeze(mask_image)
        mask_image = Image.fromarray((mask_image * 255).astype('uint8'))

        buffered = io.BytesIO()
        mask_image.save(buffered, format="PNG")
        mask_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        randomint = random.randint(1000000000, 9999999999)

        json_data = {
            'status': 'success',
            'image_path': image_path,
            'coordinates_positive': coordinates_positive,
            'coordinates_negative': coordinates_negative,
            'timestamp': timestamp,
            'random_int': randomint,
            'mask': mask_base64,
            'request_type': request_type
        }

        json_filepath = save_json_with_timestamp(json_data, randomint, request_type)

        return jsonify(json_data)

    except Exception as e:
        error_response = save_error_json(str(e), request_type)
        return jsonify(error_response), 500
    finally:
        with state.lock:
            state.is_generating = False

@app.route('/v1/models/segment:text', methods=['POST'])
def segmentation_text():
    """세그멘테이션 A 모델의 텍스트 기반 세그멘테이션"""
    state = model_states["SEGMENT_A"]

    with state.lock:
        if state.is_generating:
            return jsonify({'status': 'error', 'message': 'Segmentation engine A is busy'}), 503
        state.is_generating = True

    try:
        #time.sleep(5)
        data = request.get_json()
        #data['request_type'] = 'A'  # request_type 강제 설정
        request_type = data.get('request_type')
        base64_image = data.get('image')
        image_path = data.get('image_path', 'no image path')

        if not base64_image:
            error_response = save_error_json(str(e), request_type)
            return jsonify(error_response), 400

        image_np = load_image_base64_to_numpy(base64_image)
        image_tensor = torch.from_numpy(image_np).float()

        target = data['target']
        seed = random.randint(0, 2 ** 32 - 1)

        _, __, __, F_out = Florence2Run().encode(
            image=image_tensor,
            florence2_model=F_model_a[0],
            text_input=target,
            task="caption_to_phrase_grounding",
            fill_mask=True,
            seed=seed
        )

        Fco_out, _ = Florence2toCoordinates().segment(data=F_out, index="0")
        segmenter = Sam2Segmentation()

        mask_output = segmenter.segment(
            image=image_tensor,
            sam2_model=sam2_model_a[0],
            keep_model_loaded=True,
            coordinates_positive=Fco_out,
        )

        mask_image = mask_output[0].cpu().numpy()
        mask_image = np.squeeze(mask_image)
        mask_image = Image.fromarray((mask_image * 255).astype('uint8'))

        buffered = io.BytesIO()
        mask_image.save(buffered, format="PNG")
        mask_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        json_data = {
            'status': 'success',
            'seed': seed,
            'image_path': image_path,
            'target_text': target,
            'timestamp': timestamp,
            'mask': mask_base64,
            'request_type': request_type
        }

        json_filepath = save_json_with_timestamp(json_data, seed, request_type)


        return jsonify(json_data)

    except Exception as e:
        error_response = save_error_json(str(e), request_type)
        return jsonify(error_response), 500
    finally:
        with state.lock:
            state.is_generating = False

@app.route('/v1/models/segment:text2', methods=['POST'])
def segmentation_text2():
    """세그멘테이션 B 모델의 텍스트 기반 세그멘테이션"""
    state = model_states["SEGMENT_B"]

    with state.lock:
        if state.is_generating:
            return jsonify({'status': 'error', 'message': 'Segmentation engine B is busy'}), 503
        state.is_generating = True

    try:
        #time.sleep(5)
        data = request.get_json()
        #data['request_type'] = 'B'  # request_type 강제 설정
        request_type = data.get('request_type')
        base64_image = data.get('image')
        image_path = data.get('image_path', 'no image path')

        if not base64_image:
            error_response = save_error_json(str(e), request_type)
            return jsonify(error_response), 400

        image_np = load_image_base64_to_numpy(base64_image)
        image_tensor = torch.from_numpy(image_np).float()

        target = data['target']
        seed = random.randint(0, 2 ** 32 - 1)

        _, __, __, F_out = Florence2Run().encode(
            image=image_tensor,
            florence2_model=F_model_b[0],  # B 모델 사용
            text_input=target,
            task="caption_to_phrase_grounding",
            fill_mask=True,
            seed=seed
        )

        Fco_out, _ = Florence2toCoordinates().segment(data=F_out, index="0")
        segmenter = Sam2Segmentation()

        mask_output = segmenter.segment(
            image=image_tensor,
            sam2_model=sam2_model_b[0],  # B 모델 사용
            keep_model_loaded=True,
            coordinates_positive=Fco_out,
        )

        mask_image = mask_output[0].cpu().numpy()
        mask_image = np.squeeze(mask_image)
        mask_image = Image.fromarray((mask_image * 255).astype('uint8'))

        buffered = io.BytesIO()
        mask_image.save(buffered, format="PNG")
        mask_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        json_data = {
            'status': 'success',
            'seed': seed,
            'image_path': image_path,
            'target_text': target,
            'timestamp': timestamp,
            'mask': mask_base64,
            'request_type': request_type
        }

        json_filepath = save_json_with_timestamp(json_data, seed, request_type)


        return jsonify(json_data)

    except Exception as e:
        error_response = save_error_json(str(e), request_type)
        return jsonify(error_response), 500
    finally:
        with state.lock:
            state.is_generating = False

@app.route("/v1/models/t2i:predict", methods=["POST"])
def model_predict_a():
    """모델 A를 위한 예측 엔드포인트"""
    state = model_states["A"]

    with state.lock:
        if state.is_generating:
            return jsonify({
                'status': 'error',
                'message': 'Engine A is currently busy. Please try again later.'
            }), 503
        state.is_generating = True

    try:
        data = request.get_json()
        # request_type을 'A'로 강제 설정
        #data['request_type'] = 'A'

        response_body = generation_image_by_model("model", 'clip', 'vae')
        return jsonify(response_body)

    except Exception as e:
        error_response = save_error_json(str(e), request_type)
        return jsonify(error_response), 500
    finally:
        with state.lock:
            state.is_generating = False

@app.route("/v1/models/t2i:predict2", methods=["POST"])
def model_predict_b():
    """모델 B를 위한 예측 엔드포인트"""
    state = model_states["B"]

    with state.lock:
        if state.is_generating:
            return jsonify({
                'status': 'error',
                'message': 'Engine B is currently busy. Please try again later.'
            }), 503
        state.is_generating = True

    try:
        data = request.get_json()
        # request_type을 'B'로 강제 설정
        #data['request_type'] = 'B'

        response_body = generation_image_by_model("model_b", 'clip_b', 'vae_b')
        return jsonify(response_body)

    except Exception as e:
        error_response = save_error_json(str(e), request_type)
        return jsonify(error_response), 500
    finally:
        with state.lock:
            state.is_generating = False



class ResizeMask:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
                "width": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, "display": "number"}),
                "height": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, "display": "number"}),
                "keep_proportions": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("MASK", "INT", "INT",)
    RETURN_NAMES = ("mask", "width", "height",)
    FUNCTION = "resize"
    CATEGORY = "KJNodes/masking"

    def resize(self, mask, width, height, keep_proportions):
        if keep_proportions:
            _, oh, ow = mask.shape
            width = ow if width == 0 else width
            height = oh if height == 0 else height
            ratio = min(width / ow, height / oh)
            width = round(ow * ratio)
            height = round(oh * ratio)
        outputs = mask.unsqueeze(1)
        outputs = F.interpolate(outputs, size=(height, width), mode="nearest")
        outputs = outputs.squeeze(1)
        return (outputs, outputs.shape[2], outputs.shape[1],)


@app.route('/v1/models/keyvisual:predict', methods=['POST'])
def keyvisual_predict():

    import os
    try:
        print("File exists check:", os.path.exists("/root/app/custom_nodes/keyvisual/wcoh.py"))
    except Exception as e:
        print("Path check error:", str(e))

    try:
        from custom_nodes.keyvisual import wcoh as wcoh_module
        text_generator = wcoh_module.wcoh()
    except Exception as e:
        print("Import error:", str(e))

    import sys
    print("sys.path:", sys.path)

    """키비주얼 모델의 예측 엔드포인트"""
    state = model_states["KEYVISUAL"]


    
    with state.lock:
        if state.is_generating:
            return jsonify({
                'status': 'error',
                'message': 'Keyvisual engine is currently busy. Please try again later.'
            }), 503
        state.is_generating = True

    try:
        data = request.get_json()

        # 필수 파라미터
        positive_prompt = data.get('positive_prompt', '')
        negative_prompt = data.get('negative_prompt', '')
        visual_text = data.get('visual_text', '')
        request_type = data.get('request_type','')

        # 이미지 생성 설정 (선택적 파라미터)
        width = data.get('width', 832)
        height = data.get('height', 1152)
        steps = data.get('steps', 8)
        cfg = data.get('cfg', 4)
        sampler_name = data.get('sampler_name', 'dpmpp_sde')
        scheduler = data.get('scheduler', 'karras')
        seed = data.get('seed', random.randint(0, 2 ** 32 - 1))
        controlnet_stop = data.get('controlnet_stop', 0.8)
        controlnet_weight = data.get('controlnet_weight', 1.5)

        # Fooocus Loader로 프롬프트 업데이트
        loader = FooocusLoader()
        pipe = loader.fooocus_loader(
            positive=positive_prompt,
            negative=negative_prompt,
            base_model_name="DreamShaperXL_Turbo_v2_1.safetensors",
            vae_name="sdxl_vae.safetensors",
            refiner_model_name="None",
            refiner_switch=0.5,
            refiner_swap_method="joint",
            clip_skip=2,
            resolution=f"{width} x {height}",
            image_number=1
        )["result"][0]

        # PreKSampler로 이미지 생성 설정 업데이트
        preksampler = FooocusPreKSampler()
        result = preksampler.fooocus_preKSampler(
            pipe=pipe,
            steps=steps,
            cfg=cfg,
            sampler_name=sampler_name,
            scheduler=scheduler,
            denoise=1.0,
            settings="Simple",
            controlnet_softness=0.25,
            sharpness=2.0,
            adaptive_cfg=7.0,
            adm_scaler_positive=1.5,
            adm_scaler_negative=0.8,
            adm_scaler_end=0.3,
            freeu_enabled=False,
            seed=seed
        )

        pipe = result["result"][0]
        model = result["result"][1]

        

        # 텍스트 이미지 생성
        text_generator = keyvisual_nodes["wcoh"]()
        text_image = text_generator.text_to_image(
            text=visual_text,
            selected_font=data.get('font', 'YOnepick-Bold'),
            align=data.get('align', 'center'),
            wrap=data.get('wrap', 0),
            font_size=data.get('font_size', 160),
            color=data.get('color', 'red'),
            outline_size=data.get('outline_size', 0),
            outline_color=data.get('outline_color', 'blue'),
            margin_x=data.get('margin_x', 10),
            margin_y=data.get('margin_y', 250),
            line_spacing=data.get('line_spacing', 10),
            width=width,
            height=height,
            swap=True,
            arc_text=False,
            arc_radius=100,
            arc_start_angle=180,
            arc_end_angle=360
        )[0]

        # ControlNet 적용
        controlnet = FooocusControlnet()
        pipe, _ = controlnet.apply_controlnet(
            pipe=pipe,
            image=text_image,
            cn_type="PyraCanny",
            cn_stop=controlnet_stop,
            cn_weight=controlnet_weight,
            skip_cn_preprocess=False
        )

        # KSampler로 최종 이미지 생성
        ksampler = FooocusKsampler()
        result = ksampler.ksampler(
            pipe=pipe,
            model=model,
            image_output="Hide",
            save_prefix="ComfyUI"
        )

        # 중간 이미지 추출
        if isinstance(result, dict) and "result" in result:
            mid_img = result["result"][1]
        elif isinstance(result, tuple):
            mid_img = result[1]

        # LG U+ 로고 오버레이
        loader = LoadImage()
        logo_img, logo_mask = loader.load_image(data.get('logo1_path', '/root/app/custom_nodes/keyvisual/LG_U+_CI.svg.png'))

        # 로고 리사이즈
        image_resizer = ImageResize()
        logo_img = image_resizer.execute(
            image=logo_img,
            width=data.get('logo1_width', 128),
            height=data.get('logo1_height', 32),
            interpolation="nearest",
            method="stretch",
            condition="always",
            multiple_of=0
        )[0]
        logo_mask = ResizeMask().resize(
            mask=logo_mask,
            width=data.get('logo1_width', 128),
            height=data.get('logo1_height', 32),
            keep_proportions=False
        )[0]

        # 첫 번째 로고 오버레이
        Imgoverlay = TSC_ImageOverlay()
        ov1_image = Imgoverlay.apply_overlay_image(
            base_image=mid_img,
            overlay_image=logo_img,
            optional_mask=logo_mask,
            overlay_resize="None",
            resize_method="nearest-exact",
            rescale_factor=1.0,
            width=width,
            height=height,
            x_offset=data.get('logo1_x', 20),
            y_offset=data.get('logo1_y', 20),
            rotation=0,
            opacity=0
        )[0]

        # ixi 로고 오버레이
        ixi_img, ixi_mask = loader.load_image(data.get('logo2_path', '/root/app/custom_nodes/keyvisual/produced by ixi.png'))
        ixi_img = image_resizer.execute(
            image=ixi_img,
            width=data.get('logo2_width', 400),
            height=data.get('logo2_height', 64),
            interpolation="nearest",
            method="stretch",
            condition="always",
            multiple_of=0
        )[0]
        ixi_mask = ResizeMask().resize(
            mask=ixi_mask,
            width=data.get('logo2_width', 400),
            height=data.get('logo2_height', 64),
            keep_proportions=False
        )[0]

        ov2_image = Imgoverlay.apply_overlay_image(
            base_image=ov1_image,
            overlay_image=ixi_img,
            optional_mask=ixi_mask,
            overlay_resize="None",
            resize_method="nearest-exact",
            rescale_factor=1.0,
            width=width,
            height=height,
            x_offset=data.get('logo2_x', 230),
            y_offset=data.get('logo2_y', 1070),
            rotation=0,
            opacity=0
        )[0]

        # 최종 이미지 변환
        final_image_np = ov2_image.cpu().numpy()
        final_image_np = np.squeeze(final_image_np)
        final_image_pil = Image.fromarray((final_image_np * 255).astype('uint8'))

        # 이미지 저장
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        randomint = random.randint(1000000000, 9999999999)

        image_filepath = save_image_with_timestamp(final_image_pil, seed, request_type)


        # base64 인코딩
        buffered = io.BytesIO()
        final_image_pil.save(buffered, format="PNG")
        image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

        # 응답 데이터 생성
        json_data = {
            'status': 'success',
            'key_visual': image_base64,
            'parameters': {
                'positive_prompt': positive_prompt,
                'negative_prompt': negative_prompt,
                'visual_text': visual_text,
                'width': width,
                'height': height,
                'steps': steps,
                'cfg': cfg,
                'sampler_name': sampler_name,
                'scheduler': scheduler,
                'seed': seed,
                'controlnet_stop': controlnet_stop,
                'controlnet_weight': controlnet_weight,
                'logo1_settings': {
                    'path': data.get('logo1_path', 'keyvisual/LG_U+_CI.svg.png'),
                    'width': data.get('logo1_width', 128),
                    'height': data.get('logo1_height', 32),
                    'x': data.get('logo1_x', 20),
                    'y': data.get('logo1_y', 20)
                },
                'logo2_settings': {
                    'path': data.get('logo2_path', 'keyvisual/produced by ixi.png'),
                    'width': data.get('logo2_width', 400),
                    'height': data.get('logo2_height', 64),
                    'x': data.get('logo2_x', 230),
                    'y': data.get('logo2_y', 1070)
                }
            },
            'timestamp': timestamp,
            'request_type': request_type
        }

        # JSON 파일 저장
        json_filepath = save_json_with_timestamp(json_data, seed, request_type)

        
        return jsonify(json_data)

    except Exception as e:
        error_response = save_error_json(str(e), request_type)
        return jsonify(error_response), 500
    finally:
        with state.lock:
            state.is_generating = False

@app.route('/v1/models/welcomeposter:predict', methods=['POST'])
def welcomePoster_start():
    state = model_states["WELCOME_POSTER"]
    
    with state.lock:
        if state.is_generating:
            return jsonify({'status': 'error', 'message': 'Welcome Poster engine is busy'}), 503
        state.is_generating = True
    
    try:
        data = request.get_json()
        if data is None:
            raise ValueError("JSON 데이터가 필요합니다.")

        team_name = data.get("team_name", "즐거운유플팀")
        user_name = data.get("user_name", "김유플")
        message = data.get("message", "유플러스에서 힘찬 도약을 응원합니다.")
        concept = data.get("concept", "1")
        gender = data.get("gender", "male")
        source_image_b64 = data.get("source_image")
        request_type = data.get("request_type", "welcome_poster")

        if not source_image_b64:
            raise ValueError("source_image (base64) 데이터가 필요합니다.")

        # base64 디코딩 후 임시 파일로 저장
        decoded_image = base64.b64decode(source_image_b64)
        temp_dir = tempfile.gettempdir()
        temp_source_path = os.path.join(temp_dir, "temp_source.jpg")
        with open(temp_source_path, "wb") as f:
            f.write(decoded_image)

        # 컨셉 이미지 선택
        base_concept_folder = "/data/public/conceptImages"
        random_image_folder = os.path.join(base_concept_folder, concept, gender)
        random_image_path = WCOHRandomImage(random_image_folder).select_random_image()

        # 얼굴 교체 및 텍스트 추가
        swapped_face = global_face_changer.run(temp_source_path, random_image_path)
        swapped_face_pil = Image.fromarray(cv2.cvtColor(swapped_face, cv2.COLOR_BGR2RGB))
        final_image = WCOHTextOnImage().add_text(swapped_face_pil, team_name, user_name, message,concept)

        # 결과 이미지를 base64로 인코딩
        buffered = io.BytesIO()
        final_image.save(buffered, format="JPEG")
        encoded_image = base64.b64encode(buffered.getvalue()).decode("utf-8")

        # 임시 파일 삭제
        if os.path.exists(temp_source_path):
            os.remove(temp_source_path)

        # 응답 데이터 생성
        response_data = {
            'status': 'success',
            'image': encoded_image,
            'parameters': {
                'team_name': team_name,
                'user_name': user_name,
                'message': message,
                'concept': concept,
                'gender': gender
            },
            'request_type': request_type
        }

        # JSON 파일 저장
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        randomint = random.randint(1000000000, 9999999999)
        json_filepath = save_json_with_timestamp(response_data, randomint, request_type)

        # 이미지 저장
        image_filepath = save_image_with_timestamp(final_image, randomint, request_type)

        return jsonify(response_data)

    except Exception as e:
        error_response = save_error_json(str(e), request_type)
        return jsonify(error_response), 500
    finally:
        with state.lock:
            state.is_generating = False


#솔루션팀
try:
    name = os.environ['TLO_APP_ID']
except KeyError as e:
    name = "T2I_MODEL_APP"


@app.route(f"/v1/models", methods=["GET"])
def health():
    return "OK"

@app.route(f"/v1/models/{name}", methods=["GET"])
def mhealth():
    return "OK"

# Main execution
if __name__ == '__main__':
    print("Initializing models...")

    folder_paths.add_model_folder_path('checkpoints', '/data/public/checkpoint')
    folder_paths.add_model_folder_path('vae', '/data/public/vae')

    initialize_models("A")
    initialize_models("B")
    initialize_models("SEGMENT_A")  # 세그멘테이션 A 모델 초기화
    initialize_models("SEGMENT_B")  # 세그멘테이션 B 모델 초기화
    initialize_models("KEYVISUAL")
    initialize_models("WELCOME_POSTER")  # Welcome Poster 초기화 추가

    print("All models loaded successfully!")
    app.run(host='0.0.0.0', port=8080, debug=True, use_reloader=False, threaded=True)
else:
    print("Initializing models...")

    folder_paths.add_model_folder_path('checkpoints', '/data/public/models/checkpoint')
    folder_paths.add_model_folder_path('vae', '/data/public/models/vae')

    initialize_models("A")
    initialize_models("B")
    initialize_models("SEGMENT_A")  # 세그멘테이션 A 모델 초기화
    initialize_models("SEGMENT_B")  # 세그멘테이션 B 모델 초기화
    initialize_models("KEYVISUAL")
    initialize_models("WELCOME_POSTER")  # Welcome Poster 초기화 추가

    print("All models loaded successfully!")