# Import necessary libraries
import torch
import gradio as gr
import webbrowser
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from PIL import Image

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
    print("⚠️ Warning: Running on CPU, performance may be slow.")

# Load Text-to-Image model
print("🔄 Loading Stable Diffusion txt2img model...")
pipe_txt2img = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)
print("✅ Text-to-Image model loaded!")

# Load Image-to-Image model
print("🔄 Loading Stable Diffusion img2img model...")
pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)
print("✅ Image-to-Image model loaded!"
)

# Function to generate images from text
def generate_txt2img(prompt, steps=50, guidance=7.5, width=512, height=512, seed=-1, save_format="png"):
    generator = torch.manual_seed(seed) if seed != -1 else None
    image = pipe_txt2img(
        prompt, num_inference_steps=steps, guidance_scale=guidance, width=width, height=height,
        generator=generator
    ).images[0]
    
    output_path = f"generated_image.{save_format}"  # Save image in the requested format
    image.save(output_path, format=save_format.upper())  # Save in the selected format
    print(f"Image saved to {output_path}")
    return output_path

# Function to transform images using img2img
def generate_img2img(prompt, image, strength=0.5, steps=50, guidance=7.5, width=512, height=512, seed=-1, save_format="png"):
    generator = torch.manual_seed(seed) if seed != -1 else None
    image = pipe_img2img(
        prompt, image=image, strength=strength, num_inference_steps=steps, guidance_scale=guidance,
        width=width, height=height, generator=generator
    ).images[0]
    
    output_path = f"modified_image.{save_format}"  # Save image in the requested format
    image.save(output_path, format=save_format.upper())  # Save in the selected format
    print(f"Image saved to {output_path}")
    return output_path

# Define Gradio UI
def create_ui():
    with gr.Blocks(title="DiffuGen: AI Image Generation") as demo:
        gr.Markdown("# 🌟 DiffuGen - AI Image Generator")

        # Text-to-Image Tab
        with gr.Tab("📷 Text to Image"):
            with gr.Row():
                prompt = gr.Textbox(label="Enter a text prompt")

            with gr.Row():
                steps = gr.Slider(10, 100, value=50, step=10, label="Steps")
                guidance = gr.Slider(1, 15, value=7.5, label="Guidance Scale")

            with gr.Row():
                width = gr.Slider(256, 1024, value=512, step=64, label="Width")
                height = gr.Slider(256, 1024, value=512, step=64, label="Height")
                seed = gr.Number(value=-1, label="Seed (-1 for random)")

            with gr.Row():
                save_format = gr.Dropdown(
                    choices=["png", "jpg"], value="png", label="Select Image Format"
                )

            generate_btn = gr.Button("🚀 Generate Image")
            output_image = gr.Image(label="Generated Image", type="pil")

            generate_btn.click(
                generate_txt2img, 
                inputs=[prompt, steps, guidance, width, height, seed, save_format], 
                outputs=output_image
            )

        # Image-to-Image Tab
        with gr.Tab("🖼️ Image to Image"):
            with gr.Row():
                prompt_img2img = gr.Textbox(label="Enter a prompt")

            with gr.Row():
                input_img = gr.Image(label="Upload Image", type="pil")

            with gr.Row():
                strength = gr.Slider(0.1, 1.0, value=0.5, label="Denoising Strength")
                steps_img2img = gr.Slider(10, 100, value=50, label="Steps")
                guidance_img2img = gr.Slider(1, 15, value=7.5, label="Guidance Scale")

            with gr.Row():
                width_img2img = gr.Slider(256, 1024, value=512, step=64, label="Width")
                height_img2img = gr.Slider(256, 1024, value=512, step=64, label="Height")
                seed_img2img = gr.Number(value=-1, label="Seed (-1 for random)")

            with gr.Row():
                save_format_img2img = gr.Dropdown(
                    choices=["png", "jpg"], value="png", label="Select Image Format"
                )

            generate_img_btn = gr.Button("🔄 Transform Image")
            output_img2img = gr.Image(label="Modified Image", type="pil")

            generate_img_btn.click(
                generate_img2img,
                inputs=[prompt_img2img, input_img, strength, steps_img2img, guidance_img2img, width_img2img, height_img2img, seed_img2img, save_format_img2img],
                outputs=output_img2img
            )

    return demo

# Launch Gradio WebUI
web_ui = create_ui()
url = web_ui.launch(share=True)

# Automatically open the WebUI in a new browser tab
webbrowser.open(url)