import gradio as gr
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
from PIL import Image
import numpy as np

MODEL_ID = "timbrooks/instruct-pix2pix"
device = "cpu"
dtype = torch.float32

pipe = None

def load_model():
    global pipe
    
    if pipe is None:
        print("Loading InstructPix2Pix model...")
        try:
            pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
                MODEL_ID,
                torch_dtype=dtype,
                safety_checker=None
            )
            pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
            pipe = pipe.to(device)
            
            # Optimasi untuk CPU
            pipe.enable_attention_slicing()
            
            print("✅ Model loaded successfully!")
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            return None
    
    return pipe

def generate_image(
    input_image,
    prompt,
    image_guidance_scale,
    guidance_scale,
    num_inference_steps,
    seed
):
    if input_image is None:
        return None, "❌ Please upload an image first!"
    
    if not prompt or prompt.strip() == "":
        return None, "❌ Please enter editing instructions!"
    
    pipeline = load_model()
    
    if pipeline is None:
        return None, "❌ Failed to load model."
    
    try:
        generator = torch.Generator(device=device)
        if seed != -1:
            generator = generator.manual_seed(seed)
        
        if isinstance(input_image, np.ndarray):
            input_image = Image.fromarray(input_image)
        
        input_image = input_image.resize((512, 512), Image.Resampling.LANCZOS)
        
        print(f"🎨 Editing with: {prompt}")
        
        result = pipeline(
            prompt=prompt,
            image=input_image,
            num_inference_steps=num_inference_steps,
            image_guidance_scale=image_guidance_scale,
            guidance_scale=guidance_scale,
            generator=generator
        )
        
        output_image = result.images[0]
        
        success_msg = f"✅ Edit complete!\n📝 Instruction: {prompt}"
        
        return output_image, success_msg
        
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        print(error_msg)
        return None, error_msg

custom_css = """
.main-header {
    text-align: left;
    text-size: 20px;
    color: white;
    margin-top: 1rem;
    margin-bottom: 2rem;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    <div class="main-header">
        <h1>🎨 InstructPix2Pix Image Editor</h1>
        <p>Edit images with natural language instructions | CPU Optimized</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="📤 Upload Image",
                type="pil"
            )
            
            prompt = gr.Textbox(
                label="✏️ Edit Instructions",
                placeholder="Examples:\n- Make it sunset\n- Turn the cat into a dog\n- Add snow\n- Make it look like a painting",
                lines=3
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                image_guidance_scale = gr.Slider(
                    minimum=1.0,
                    maximum=2.0,
                    value=1.5,
                    step=0.1,
                    label="Image Guidance (how much to preserve original)"
                )
                
                guidance_scale = gr.Slider(
                    minimum=5.0,
                    maximum=15.0,
                    value=7.5,
                    step=0.5,
                    label="Text Guidance (prompt strength)"
                )
                
                num_inference_steps = gr.Slider(
                    minimum=10,
                    maximum=50,
                    value=20,
                    step=5,
                    label="Steps (quality vs speed)"
                )
                
                seed = gr.Number(
                    value=-1,
                    label="Seed (-1 for random)",
                    precision=0
                )
            
            generate_btn = gr.Button("🚀 Edit Image", variant="primary", size="lg")
            
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="✨ Edited Result",
                type="pil"
            )
            
            status_text = gr.Textbox(
                label="Status",
                lines=3,
                interactive=False
            )
            
    generate_btn.click(
        fn=generate_image,
        inputs=[
            input_image,
            prompt,
            image_guidance_scale,
            guidance_scale,
            num_inference_steps,
            seed
        ],
        outputs=[output_image, status_text]
    )

if __name__ == "__main__":
    demo.queue(max_size=5)
    demo.launch(server_name="0.0.0.0", server_port=7860)