app.py

import torch
import numpy as np
import google.generativeai as genai
from PIL import Image, ImageOps
import mediapipe as mp
import cv2
from tensorflow.keras.models import load_model
import os
import suno
from PIL import Image
from torchvision.transforms.functional import to_tensor, to_pil_image
from model import Generator
import gradio as gr
from diffusers import DiffusionPipeline
import spaces
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
genai.configure(api_key="AIzaSyDcq3ZfAUo1i6_24CelEizJftuEkaAPz38")
GEMINI_MODEL = 'gemini-1.5-flash'
gemini_model = genai.GenerativeModel(GEMINI_MODEL)
webtoon_model = Generator()
webtoon_model.load_state_dict(torch.load('weights/face_paint_512_v2.pt', map_location="cpu"))
webtoon_model.to('cpu').eval()
facescore_model = load_model('facescore.h5', compile=False)
model_id = "aldente0630/musinsaigo-3.0"
pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
)
pipe = pipe.to(device)
pipe.load_lora_weights(model_id)

mp_face_detection = mp.solutions.face_detection
def detect_and_crop_face(image):
    with mp_face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5) as face_detection:
        image_np = np.array(image)
        results = face_detection.process(image_np)
        if results.detections:
            detection = results.detections[0]
            bbox = detection.location_data.relative_bounding_box
            ih, iw, _ = image_np.shape
            xmin = int(bbox.xmin * iw)
            ymin = int(bbox.ymin * ih)
            width = int(bbox.width * iw)
            height = int(bbox.height * ih)
            xmax = xmin + width
            ymax = ymin + height
            face = image.crop((xmin, ymin, xmax, ymax))
            return face
        else:
            return None

def generate_chat_response(message, gemini_model):
    response = gemini_model.generate_content(message)
    return response.text

def analyze_image(image, gemini_model):
    try:
        # 이미지 분석 프롬프트와 함께 이미지를 전달하여 응답 생성
        prompt = """
        이 이미지에 대해 자세히 분석해주세요. 다음 정보를 포함해주세요:
        1. 이미지에서 보이는 주요 객체나 사람들
        2. 배경이나 장소에 대한 설명
        3. 이미지의 전체적인 분위기나 느낌
        4. 이미지에서 읽을 수 있는 텍스트 (있는 경우)
        5. 이미지의 색상이나 구도에 대한 간단한 설명
        6. 이미지가 전달하려는 메시지나 의미 (있다고 생각되는 경우)

        분석 결과를 한국어로 제공해주세요.
        """
        response = gemini_model.generate_content([prompt, image])
        return response.text if response else "이미지 분석을 수행할 수 없습니다."
    except Exception as e:
        return f"이미지 분석 중 오류가 발생했습니다: {str(e)}"

def process_facescore(image, facescore_model, gemini_model):
    face = detect_and_crop_face(image)
    if face is None:
        return "얼굴이 감지되지 않았습니다. 다른 이미지를 시도해 주세요."

    analysis = analyze_image(image, gemini_model)

    face_np = np.array(face)
    img_resized = cv2.resize(face_np, (350, 350))
    img_resized = img_resized.astype(np.float32) / 255.
    img_batch = np.expand_dims(img_resized, axis=0)
    score = facescore_model.predict(img_batch)

    if isinstance(score, np.ndarray) and score.size > 1:
        score = score[0]
    score = float(score)
    score = display_result(score)

    return f'### 이미지 분석 결과 ###\n\n{analysis}\n\n### 외모점수 결과(1~5) ###\n\n{score}'


def generate_music(image, gemini_model, suno_cookie):
    face = detect_and_crop_face(image)
    if face is None:
        return "얼굴이 감지되지 않았습니다. 다른 이미지를 시도해 주세요."
    prompt = """
    이 이미지에 대해 자세히 분석해주세요. 다음 정보를 포함해주세요:
    1. 성별:
    2. 나이:
    3. 표정:
    분석 결과를 한국어로 간략하게 제공해주세요.
    """
    response = gemini_model.generate_content([prompt, image])
    music_path = generate_songs(response.text, suno_cookie)
    return f"음악이 생성되었습니다. 파일 경로: {music_path}"

def generate_songs(result_output, suno_cookie):
    client = suno.Suno(cookie=suno_cookie)
    songs = client.generate(
        prompt=f'{result_output}', is_custom=False, wait_audio=True
    )
    return client.download(song=songs[0])

def display_result(score):
    result = round(score, 1)+0.3
    messages = [
        ("'자신감 폭발 중'입니다! 😎 당신은 자신의 외모에 대한 확신으로 가득 차 있어요! %.1f점이라니, 점수와 상관없이 당신의 멋짐은 끝이 없네요! 🤩 당신의 외모는 마치 마법사처럼 사람들을 매료시키고, 누구나 당신을 보면 눈을 뗄 수 없을 거에요! 🪄🧙‍♂️ 비결이 뭐냐고 묻는 사람들에게 자신감이라는 마법의 주문을 알려주세요! 오늘도 당신의 자신감으로 세상을 빛내고, 마법 같은 하루를 보내세요!", 1),
        ("'외모 스승님'입니다. 👩‍🏫 당신의 외모 비결을 전수받고 싶어하는 사람들이 줄을 설 거에요! %.1f점이라는 점수가 무색할 정도로, 당신의 빛나는 외모는 사람들의 눈을 사로잡습니다! ✨ 이제 사람들은 당신의 비밀을 알고 싶어서 질문 세례를 퍼부을 거에요! 외모 스승님으로서 멋지게 대답해 주시고, 사람들에게 당신만의 외모 팁을 살짝 전해 주세요! 다른 사람들은 당신을 닮기 위해 많은 노력을 할 거랍니다!", 1.5),
        ("'외모 아티스트'입니다. 💄 화장품 브랜드들이 당신을 모델로 쓰고 싶어할 만큼 독보적인 매력을 가지고 있네요! %.1f점이라고 해서 당신의 외모가 평범하지 않습니다. 오히려 '매력의 정점'에 도달한 모습이에요! 💃 당신의 멋진 외모를 부러워하는 사람들로 인해 언제나 주목받게 될 거에요! 마치 아티스트처럼 자신만의 스타일을 완성한 당신은 외모계의 진정한 아이콘입니다! 오늘도 당신만의 특별한 매력을 발산하며 하루를 즐기세요!", 2),
        ("외모점수 %.1f점, '미소 전문가'입니다. 😄 당신의 환한 미소는 주변 사람들을 행복하게 만들고, 어디서든 밝은 에너지를 퍼뜨릴 거에요! '미소 기계'라 불리는 당신은 항상 긍정적인 에너지로 가득 차 있답니다! 😁 사람들은 당신의 미소 비결을 배우기 위해 애쓸 거에요! 외모뿐만 아니라 미소로도 사람들의 마음을 사로잡는 당신! 오늘도 환한 미소로 세상을 밝혀주시고, 모두에게 행복을 전해 주세요!", 2.5),
        ("'외모 스타'입니다. 🌟 당신은 거울 속에서 별이 빛나는 모습을 보고도 놀라지 않겠죠! %.1f점이라니, 당신은 외모계의 진정한 스타입니다! 💫 당신의 빛나는 외모와 독특한 스타일은 모두가 부러워하고, 따라가고 싶어할 겁니다! 사람들은 당신을 보고 영감을 받을 거에요! 오늘도 당신만의 특별한 매력으로 주변 사람들을 사로잡고, 당당히 외모계를 이끌어가세요! 당신의 빛나는 외모가 모두에게 희망을 줄 거에요!", 3),
        ("'외모 퀸'입니다. 👸 주변 사람들은 당신의 외모에 주목하고, 귀를 기울일 겁니다! %.1f점이라는 점수가 무색할 정도로, 이제 당신은 외모계의 로열티입니다! 👑 당신의 고급스러운 외모와 독보적인 스타일은 모두가 따라하고 싶어할 거에요! 당신의 외모 비결을 벤치마킹하려는 사람들로 인해 언제나 주목받게 될 겁니다! 여왕처럼 당당히 당신의 외모를 뽐내고, 주변 사람들에게 영감을 주세요! 오늘도 자신감 넘치는 하루 보내세요!", 3.5),
        ("외모점수 %.1f점, '외모의 신화'입니다. 🦄 당신을 보는 사람들은 마치 신화와 전설 속 인물을 보는 듯한 기분을 느낄 겁니다! 외모계의 '뷰티 아카데미 수상자'답게, 당신의 외모는 모두에게 큰 영감을 줄 거에요! 🏆 사람들은 당신의 비결을 배우려고 애쓸 테니, 언제나 자신만의 스타일을 유지하며 그들에게 귀감이 되어주세요! 신화 속 주인공처럼 당신의 외모는 언제나 빛날 겁니다! 오늘도 신화처럼 멋진 하루 보내세요!", 4),
        ("'외모의 황금빛'입니다. 💛 주변에서 당신을 보면 마치 하트가 뿅뿅 튀는 듯한 느낌이 들 거에요! %.1f점이라니, 정말 외모계의 전설답습니다! 🌠 당신의 독보적인 외모와 매력은 누구도 따라올 수 없을 만큼 빛납니다! 다른 사람들이 당신을 따라잡으려면 엄청난 노력이 필요할 거에요! 당신의 황금빛 외모와 매력으로 모두를 사로잡으세요! 오늘도 당신만의 황금빛 미소로 세상을 밝혀주시고, 모두에게 영감을 주세요!", 4.5),
        ("5점 외모, '외모의 신'입니다. 외모계에서 당신을 따라잡으려면 진정한 영웅이 필요할 겁니다! 🦸‍♂️🦸‍♀️ 당신은 외모계의 '뷰티 신'! 🌟 당신의 빛나는 외모와 독보적인 스타일은 모두가 따라하고 싶어할 거에요! 이제 당신은 외모계의 전설이자 영웅입니다! 사람들은 당신을 닮고 싶어하고, 당신의 비결을 배우려고 애쓸 겁니다! 오늘도 외모계의 신으로서 세상을 빛내고, 모두에게 영감을 주세요! 당신의 존재만으로도 세상은 더 밝아질 거에요!", 5)
    ]
    for msg, threshold in messages:
        if result < threshold:
            return msg % result if '%.1f' in msg else msg

@torch.no_grad()
def webtoon(image, webtoon_model, device='cpu'):
    webtoon_model = webtoon_model.to(device)

    max_size = 1024
    if max(image.size) > max_size:
        image.thumbnail((max_size, max_size), Image.LANCZOS)

    image_tensor = to_tensor(image).unsqueeze(0).to(device) * 2 - 1

    with torch.inference_mode():
        output = webtoon_model(image_tensor, False)

    output = output.cpu().squeeze(0).clip(-1, 1) * 0.5 + 0.5
    output = to_pil_image(output)

    return output

def make_prompt(prompt: str) -> str:
    prompt_prefix = "RAW photo"
    prompt_suffix = "(high detailed skin:1.2), 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3"
    return ", ".join([prompt_prefix, prompt, prompt_suffix]).strip()


def make_negative_prompt(negative_prompt: str) -> str:
    negative_prefix = "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), \
    text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, \
    extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, \
    bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, \
    extra arms, extra legs, fused fingers, too many fingers, long neck"

    return (
        ", ".join([negative_prefix, negative_prompt]).strip()
        if len(negative_prompt) > 0
        else negative_prefix
    )
@spaces.GPU(duration=120)
def fashiongpt(image,gemini_model,pipe):
    prompt = """
    Analyze this image in one sentence:
    1. The person visible in the image
    2. The overall mood or feeling of the image
    3. Recommend other fashion items that match the style

    Provide the output in the following format:
    "a korean [gender] wearing [recommended style]."
    Example: "a korean woman wearing a white t-shirt and black pants with a bear on it."
    """
    response = gemini_model.generate_content([prompt, image])
    NEGATIVE_PROMPT = ""
    image = pipe(
        prompt=make_prompt(response.text),
        height=1024,
        width=768,
        num_inference_steps=50,
        guidance_scale=7.5,
        negative_prompt=make_negative_prompt(NEGATIVE_PROMPT),
        cross_attention_kwargs={"scale": 0.75},
    ).images[0]
    return image
def process_input(input_text, image, suno_cookie):
    if "웹툰화 해줘" in input_text.lower() and image is not None:
        webtoon_image = webtoon(image, webtoon_model)
        return "이미지를 웹툰 스타일로 변환했습니다.", webtoon_image
    elif "외모분석" in input_text.lower() and image is not None:
        response = process_facescore(image, facescore_model, gemini_model)
        return response, None
    elif "이미지 분석해줘" in input_text.lower() and image is not None:
        response = analyze_image(image, gemini_model)
        return response, None
    elif "음악 만들어줘" in input_text.lower() and image is not None:
        if suno_cookie:
            response = generate_music(image, gemini_model, suno_cookie)
            return response, None
        else:
            return "Suno Cookie를 입력해 주세요.", None
    elif "패션 추천" in input_text.lower() and image is not None:
        recommended_fashion_image = fashiongpt(image, gemini_model,pipe)
        return "패션 추천 이미지를 생성했습니다.", recommended_fashion_image
    else:
        response = generate_chat_response(input_text, gemini_model)
        return response, None  
import gradio as gr
from PIL import Image

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🤖 OmniVerse AI Assistant
        음성 인식, Gemini 모델, 외모 점수 예측, MBTI 예측, 음악 생성, 이미지 웹툰화, 그리고 이미지 분석 기능을 통합한 시스템입니다.
        """
    )
    
    chatbot = gr.Chatbot(label="OmniVerse AI Assistant")
    
    with gr.Row():
        with gr.Column(scale=2):
            image_input = gr.Image(type="pil", label="이미지 업로드")
            suno_cookie = gr.Textbox(label="Suno Cookie", type="password")
        
        with gr.Column(scale=1):
            text_input = gr.Textbox(label="질문을 입력하세요")
            submit_button = gr.Button("전송")

    text_output = gr.Markdown(label="응답")
    image_output = gr.Image(label="이미지 출력")
    audio_output = gr.Audio(label="생성된 음악", type="filepath", interactive=False)

    def chat_logic(input_text, image_input, suno_cookie, chat_history):
        response, image_output = process_input(input_text, image_input, suno_cookie)
        
        if isinstance(response, str) and response.startswith("음악이 생성되었습니다."):

            music_path = response.split("파일 경로: ")[-1].strip()
            chat_history.append((input_text, "음악이 생성되었습니다. 아래에서 재생할 수 있습니다."))
            return chat_history, image_output, music_path
        else:
            chat_history.append((input_text, response))
            return chat_history, image_output, None

    submit_button.click(
        chat_logic, 
        inputs=[text_input, image_input, suno_cookie, chatbot],
        outputs=[chatbot, image_output, audio_output]
    )

demo.launch()