import gradio as gr
import time
import logging

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def load_model(model_path):
    logger.info(f"Attempting to load model from {model_path}...")
    try:
        demo = gr.load(model_path)
        logger.info(f"Successfully loaded demo from {model_path}.")
        return demo
    except Exception as e:
        logger.error(f"Failed to load demo from {model_path}: {e}")
        return None

# Load the demo
demo = load_model("models/John6666/only-porn-real-porn-sdxl")

def generate_image(prompt, max_retries=5):
    # if demo is None:
    #     return None, "Error: Demo failed to load."

    for attempt in range(max_retries):
        try:
            logger.info(f"Attempt {attempt + 1}: Starting image generation with prompt: {prompt}")
            start_time = time.time()
            result = demo(prompt)
            generation_time = time.time() - start_time
            logger.info(f"Generation completed in {generation_time:.2f} seconds")
            logger.info(result)
            return result, f"Image generated successfully\nGeneration time: {generation_time:.2f} seconds"
        except Exception as e:
            logger.error(f"Error on attempt {attempt + 1}: {str(e)}")
            if attempt < max_retries - 1:
                logger.info(f"Retrying... ({attempt + 2}/{max_retries})")
            else:
                logger.error("Max retries reached. Giving up.")
                return None, f"Failed to generate image after {max_retries} attempts. Last error: {str(e)}"
    
    # This line should never be reached, but just in case:
    return None, "Unexpected error occurred"

with gr.Blocks() as demo:
    gr.Markdown("# Text-to-Image Generation")
    gr.Markdown("Enter a prompt to generate an image. The system will automatically retry up to 5 times if an error occurs.")
    
    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(label="Enter your prompt")
            generate_btn = gr.Button("Generate")
            image_output = gr.Image(label="Generated Image")
            info_output = gr.Textbox(label="Output Information", lines=5)
        
        with gr.Column(scale=1):
            file_explorer = gr.FileExplorer(label="File System")
    
    generate_btn.click(
        generate_image,
        inputs=input_text,
        outputs=[image_output, info_output]
    )

logger.info("Launching the Gradio interface...")
demo.launch()