import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
import logging
from PIL import Image
import base64
from io import BytesIO

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

logger.info("Starting model initialization...")

try:
    pipe = StableDiffusionPipeline.from_pretrained(
        "stabilityai/flux-1-schnell",
        torch_dtype=torch.float32,
        safety_checker=None,
        requires_safety_checker=False
    )
    pipe.enable_model_cpu_offload()
    logger.info("Model initialized successfully")
except Exception as e:
    logger.error(f"Error during model initialization: {str(e)}")
    raise

def generate_image(prompt, negative_prompt="", style=""):
    try:
        # Enhance prompt based on style
        if style:
            style_prompts = {
                "Realistic": "highly detailed, photorealistic, 8k resolution",
                "Artistic": "artistic style, vibrant colors, creative composition",
                "Abstract": "abstract art style, modern, conceptual",
                "Minimalist": "minimalist style, clean lines, simple composition"
            }
            prompt = f"{prompt}, {style_prompts.get(style, '')}"

        # Generate image
        logger.info(f"Generating image with prompt: {prompt}")
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=20,
            guidance_scale=7.5
        ).images[0]

        # Convert to base64
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        return {"image": f"data:image/png;base64,{img_str}"}
    except Exception as e:
        logger.error(f"Error generating image: {str(e)}")
        return {"error": str(e)}

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe the wallpaper you want to generate..."),
        gr.Textbox(label="Negative Prompt", placeholder="What you don't want to see in the image..."),
        gr.Dropdown(choices=["Realistic", "Artistic", "Abstract", "Minimalist"], label="Style", value="Realistic")
    ],
    outputs=gr.Image(type="pil"),
    title="FLUX AI Wallpaper Generator",
    description="Generate unique wallpapers using FLUX AI. Enter a prompt and select a style to create your custom wallpaper."
)

# Launch with basic settings
if __name__ == "__main__":
    try:
        logger.info("Starting Gradio interface...")
        iface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False
        )
    except Exception as e:
        logger.error(f"Error launching Gradio interface: {str(e)}")
        raise