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/the-araminta-fv1-sdxl")
#demo = load_model("models/John6666/the-araminta-experiment-fv2-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"

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt"),
    outputs=[
        gr.Image(label="Generated Image"),
        gr.Textbox(label="Output Information", lines=5)
    ],
    title="Text-to-Image Generation",
    description="Enter a prompt to generate an image. The system will automatically retry up to 5 times if an error occurs."
)

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