import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionInpaintPipeline, StableDiffusionUpscalePipeline
import numpy as np

def process_image(image, prompt, mode, scale_factor=2):
    if mode == "upscale":
        # Upscale pipeline
        pipeline = StableDiffusionUpscalePipeline.from_pretrained(
            "stabilityai/stable-diffusion-x4-upscaler"
        )
        pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
        
        # Process image
        upscaled_image = pipeline(
            prompt=prompt,
            image=image,
            noise_level=20,
            num_inference_steps=20
        ).images[0]
        
        return upscaled_image
    
    elif mode == "inpaint":
        # Inpainting pipeline
        pipeline = StableDiffusionInpaintPipeline.from_pretrained(
            "runwayml/stable-diffusion-inpainting"
        )
        pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
        
        # Create mask for extending the image
        width, height = image.size
        mask = Image.new('RGB', (width, height), 'white')
        
        # Process image
        result = pipeline(
            prompt=prompt,
            image=image,
            mask_image=mask,
            num_inference_steps=20
        ).images[0]
        
        return result

# Gradio Interface
def create_interface():
    with gr.Blocks(title="AI Image Enhancement") as interface:
        gr.Markdown("# AI Image Enhancement Studio")
        gr.Markdown("Enhance, upscale, and recreate images using AI")
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Upload Image")
                prompt = gr.Textbox(label="Prompt", placeholder="Describe the desired enhancement...")
                mode = gr.Radio(
                    choices=["upscale", "inpaint"],
                    label="Processing Mode",
                    value="upscale"
                )
                scale_factor = gr.Slider(
                    minimum=2,
                    maximum=8,
                    step=2,
                    label="Upscale Factor",
                    value=2
                )
                process_btn = gr.Button("Process Image")
            
            with gr.Column():
                output_image = gr.Image(type="pil", label="Enhanced Result")
        
        process_btn.click(
            fn=process_image,
            inputs=[input_image, prompt, mode, scale_factor],
            outputs=output_image
        )
    
    return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch()