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

# Load the Stable Diffusion img2img pipeline
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to(device)
#pipe.enable_attention_slicing()

# Predefined prompts
PROMPTS = {
    "Photorealistic": "A heavily corroded metal pipeline stretching across the ocean floor, covered in algae and barnacles, deep underwater with ambient blue lighting and floating particles, photorealistic",
    "Cinematic": "An old rusted pipeline submerged in the sea, encrusted with marine growth and decay, surrounded by dark water and shafts of light from the surface, cinematic, moody atmosphere"
}

# Inference function
def generate_image(init_image, prompt_choice, strength, guidance_scale):
    # Resize and convert the input image
    init_image = init_image.convert("RGB").resize((768, 512))
    
    # Get the selected prompt
    prompt = PROMPTS[prompt_choice]

    # Run the pipeline
    result = pipe(
        prompt=prompt,
        image=init_image,
        strength=strength,
        guidance_scale=guidance_scale
    ).images[0]

    return result

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("Corroded Pipeline Generator - Underwater Img2Img")
    with gr.Row():
        with gr.Column():
            init_image = gr.Image(label="Upload Initial Image", type="pil")
            prompt_choice = gr.Radio(choices=list(PROMPTS.keys()), label="Select Prompt", value="Photorealistic")
            strength = gr.Slider(minimum=0.2, maximum=1.0, value=0.75, step=0.05, label="Transformation Strength")
            guidance_scale = gr.Slider(minimum=1, maximum=15, value=7.5, step=0.5, label="Prompt Guidance Scale")
            generate_btn = gr.Button("Generate")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")

    generate_btn.click(fn=generate_image, inputs=[init_image, prompt_choice, strength, guidance_scale], outputs=output_image)

# Launch the app
demo.launch()