import gradio as gr
import torch
from diffusers import FluxPipeline
import spaces
import random
import gc
import logging
from io import BytesIO
from PIL import Image
import time

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

"""
This application uses the Flux.1 Lite model:
@article{flux1-lite,
  title={Flux.1 Lite: Distilling Flux1.dev for Efficient Text-to-Image Generation},
  author={Daniel Verdú, Javier Martín},
  email={dverdu@freepik.com, javier.martin@freepik.com},
  year={2024},
}
"""

@spaces.GPU(duration=70)
def initialize_model():
    logger.info("Initializing Flux.1 Lite model")
    model_id = "Freepik/flux.1-lite-8B"
    try:
        pipe = FluxPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16
        ).to("cuda")
        logger.info("Model loaded successfully")
        return pipe
    except Exception as e:
        logger.error(f"Failed to initialize model: {str(e)}")
        raise e

@spaces.GPU(duration=70)
def generate_image(
    prompt,
    guidance_scale=3.5,
    num_inference_steps=28,
    width=512,
    height=512
):
    logger.info(f"Generating image with prompt: {prompt}, guidance_scale: {guidance_scale}, steps: {num_inference_steps}, size: {width}x{height}")
    max_retries = 3
    retry_delay = 5  # seconds
    for attempt in range(max_retries):
        try:
            # Initialize model
            pipe = initialize_model()
            
            # Generate random seed
            seed = random.randint(1, 1000000)
            logger.info(f"Using seed: {seed}")
            
            # Generate image
            with torch.inference_mode():
                image = pipe(
                    prompt=prompt,
                    generator=torch.Generator(device="cpu").manual_seed(seed),
                    num_inference_steps=num_inference_steps,
                    guidance_scale=guidance_scale,
                    height=height,
                    width=width,
                ).images[0]
            
            # Convert PIL image to byte stream
            buffer = BytesIO()
            image.save(buffer, format="PNG", optimize=True)
            image_bytes = buffer.getvalue()
            buffer.close()
            
            # Reopen as PIL Image for Gradio
            image = Image.open(BytesIO(image_bytes))
            
            # Clean up
            del pipe
            torch.cuda.empty_cache()
            gc.collect()
            logger.info("Image generated and memory cleaned")
            
            return image
        except Exception as e:
            logger.error(f"Attempt {attempt+1}/{max_retries} failed: {str(e)}")
            if attempt < max_retries - 1:
                logger.info(f"Retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
                continue
            raise e
        finally:
            torch.cuda.empty_cache()
            gc.collect()

# Create the Gradio interface
demo = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(
            label="Prompt",
            placeholder="Enter your image description here...",
            value="a glass cup with beer, inside the beer a scuba diver, with a beautiful sunset background"
        ),
        gr.Slider(
            minimum=2.0,
            maximum=5.0,
            value=3.5,
            label="Guidance Scale",
            step=0.1
        ),
        gr.Slider(
            minimum=20,
            maximum=32,
            value=28,
            label="Number of Inference Steps",
            step=1
        ),
        gr.Slider(
            minimum=128,
            maximum=1024,
            value=512,
            label="Width",
            step=64
        ),
        gr.Slider(
            minimum=128,
            maximum=1024,
            value=512,
            label="Height",
            step=64
        )
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Freepik Flux.1-lite-8B Model (Zero-GPU)",
    description="Generate images using Freepik's updated Flux.1-lite-8B model with Zero-GPU allocation. Optimized with recommended guidance scale (2.0-5.0) and inference steps (20-32).",
    examples=[
        ["A close-up image of a green alien with fluorescent skin in the middle of a dark purple forest", 3.5, 28, 512, 512],
        ["a glass cup with beer, inside the beer a scuba diver, with a beautiful sunset background", 3.5, 28, 512, 512]
    ],
    cache_examples=False,
    allow_flagging="never"
)

# Launch the app
if __name__ == "__main__":
    logger.info("Launching Gradio app")
    demo.launch(server_name="0.0.0.0", server_port=7860)