import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from PIL import Image

# Device configuration
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load Stable Diffusion pipelines
text_to_image_pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)

image_to_image_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)

# Function for Text-to-Image
def text_to_image(prompt, negative_prompt, guidance_scale, num_inference_steps):
    image = text_to_image_pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
    ).images[0]
    return image


# Function for Image-to-Image
def image_to_image(prompt, negative_prompt, init_image, strength, guidance_scale, num_inference_steps):
    init_image = init_image.convert("RGB").resize((512, 512))  # Ensure the image is resized
    image = image_to_image_pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        init_image=init_image,
        strength=strength,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
    ).images[0]
    return image


# Gradio Interface
with gr.Blocks(theme='NoCrypt/miku') as demo:
    gr.Markdown("# Text-to-Image and Image-to-Image generation")
    
    with gr.Tab("Text-to-Image"):
        gr.Markdown("Generate images from text prompts")
        with gr.Row():
            text_prompt = gr.Textbox(label="Prompt", placeholder="Enter your text here...")
            text_negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter what to avoid...")
        with gr.Row():
            guidance_scale = gr.Slider(1, 20, value=7.5, step=0.1, label="Guidance Scale")
            num_inference_steps = gr.Slider(10, 100, value=50, step=1, label="Inference Steps")
        with gr.Row():
            generate_btn = gr.Button("Generate", elem_classes=["primary-button"])
        with gr.Row():
            text_output = gr.Image(label="Generated Image")

        generate_btn.click(
            text_to_image,
            inputs=[text_prompt, text_negative_prompt, guidance_scale, num_inference_steps],
            outputs=text_output,
        )

    with gr.Tab("Image-to-Image"):
        gr.Markdown(
            "Modify images - Upload an image, provide a prompt describing the transformation, and adjust settings for desired results."
        )
        with gr.Row():
            init_image = gr.Image(type="pil", label="Upload Initial Image")
        with gr.Row():
            img_prompt = gr.Textbox(label="Prompt", placeholder="Describe modifications...")
            img_negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter what to avoid...")
        with gr.Row():
            strength = gr.Slider(0.1, 1.0, value=0.75, step=0.05, label="Strength")
            img_guidance_scale = gr.Slider(1, 20, value=7.5, step=0.1, label="Guidance Scale")
            img_num_inference_steps = gr.Slider(10, 100, value=50, step=1, label="Inference Steps")
        with gr.Row():
            img_generate_btn = gr.Button("Generate", elem_classes=["primary-button"])
        with gr.Row():
            img_output = gr.Image(label="Modified Image")

        img_generate_btn.click(
            image_to_image,
            inputs=[img_prompt, img_negative_prompt, init_image, strength, img_guidance_scale, img_num_inference_steps],
            outputs=img_output,
        )

demo.launch(share=True)