import gradio as gr
from gradio_client import Client
import os
import logging

# 로깅 설정
logging.basicConfig(level=logging.INFO)

# API 클라이언트 설정
api_client = Client("http://211.233.58.202:7960/")

def respond(message):
    logging.info("Received message: %s", message)
    
    try:
        # 이미지 생성 요청
        result = api_client.predict(
            prompt=message,
            seed=123,
            randomize_seed=False,
            width=1024,
            height=576,
            guidance_scale=5,
            num_inference_steps=28,
            api_name="/infer_t2i"
        )
        logging.info("API response received: %s", result)
        
        # 결과 확인 및 처리
        if isinstance(result, dict) and 'url' in result:
            return result['url']
        elif isinstance(result, tuple):
            logging.error("Unexpected tuple response: %s", result)
            return result[0]
        else:
            raise ValueError("Unexpected API response format")
    except Exception as e:
        logging.error("Error during API request: %s", str(e))
        return "Failed to generate image due to an error."

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

# 이미지 생성을 위한 예제 프롬프트
examples = [
    ["A futuristic cityscape at sunset."],
    ["A portrait of a cat wearing a monocle."],
    ["A serene landscape with mountains in the background and a clear lake in the foreground."],
    ["A street scene from Tokyo at night, vibrant and full of lights."],
    ["An astronaut riding a horse on Mars."],
    ["A surreal painting of a tree growing books."],
    ["A cottage in a snowy forest, lit by warm lights."],
    ["A still life of various fruits and a wine glass on a table."],
    ["A digital artwork of a neon-lit alley in a cyberpunk city."],
    ["A fantasy map of a fictional world, with detailed terrain and cities."]
]

# Gradio 인터페이스 설정
demo = gr.Interface(
    fn=respond,
    inputs=gr.Textbox(label="Enter your prompt for image generation"),
    outputs=gr.Image(label="Generated Image"),
    examples=examples,
    theme="Nymbo/Nymbo_Theme",
    css=css
)

if __name__ == "__main__":
    demo.launch()