import gradio as gr
import torch
from diffusers import DiffusionPipeline  # Note: Change `FluxPipeline` to `DiffusionPipeline` if `FluxPipeline` is not correct
from PIL import Image

# Function to determine the device and handle model loading
def setup_pipeline():
    # Check for CUDA availability
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Load the diffusion model
    try:
        pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
        if device == "cpu":
            # If using CPU, ensure model is offloaded to avoid GPU-specific features
            pipeline.enable_model_cpu_offload()
        else:
            # Move model to GPU
            pipeline.to(device)
    except Exception as e:
        print(f"Error loading model: {e}")
        raise e

    return pipeline, device

pipeline, device = setup_pipeline()

def generate_image(prompt, guidance_scale=7.5, num_inference_steps=50):
    # Generate an image based on the prompt
    with torch.no_grad():
        try:
            images = pipeline(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
        except Exception as e:
            print(f"Error generating image: {e}")
            raise e
    
    # Assuming pipeline returns a list of images, just take the first one
    img = images[0]
    
    # Convert PIL image to format suitable for Gradio
    return img

# Set up Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Text to Image Generation")
    
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here...")
        guidance_scale = gr.Slider(minimum=1, maximum=15, step=0.1, value=7.5, label="Guidance Scale")
        num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, value=50, label="Number of Inference Steps")
    
    with gr.Row():
        generate_button = gr.Button("Generate Image")
    
    result = gr.Image(label="Generated Image")
    
    # Connect the function to the button
    generate_button.click(
        fn=generate_image,
        inputs=[prompt, guidance_scale, num_inference_steps],
        outputs=result
    )

# Launch the app
demo.launch()