app.py

import gradio as gr
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline
from PIL import Image
import spaces

# Global model variable
pipe = None

def load_model():
    """Load the image editing model"""
    global pipe
    if pipe is None:
        print("Loading model...")
        pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
            "timbrooks/instruct-pix2pix",
            torch_dtype=torch.float16,
            safety_checker=None
        )
        pipe.to("cuda")
        pipe.enable_attention_slicing()
        print("Model loaded successfully!")
    return pipe

@spaces.GPU
def edit_image(input_image, edit_prompt, num_steps, guidance_scale, image_guidance_scale):
    """Edit an image based on text instructions"""
    
    if input_image is None:
        return None, "❌ Please upload an image first!"
    
    if not edit_prompt or edit_prompt.strip() == "":
        return None, "❌ Please provide edit instructions!"
    
    try:
        print(f"Starting edit with prompt: {edit_prompt}")
        
        # Load model
        model = load_model()
        
        # Resize if too large
        max_size = 512
        if max(input_image.size) > max_size:
            ratio = max_size / max(input_image.size)
            new_size = tuple(int(dim * ratio) for dim in input_image.size)
            input_image = input_image.resize(new_size, Image.Resampling.LANCZOS)
        
        # Ensure dimensions are multiples of 8
        width = (input_image.width // 8) * 8
        height = (input_image.height // 8) * 8
        input_image = input_image.resize((width, height))
        
        print(f"Processing image: {width}x{height}")
        
        # Generate edited image
        result = model(
            edit_prompt,
            image=input_image,
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale,
            image_guidance_scale=image_guidance_scale,
        ).images[0]
        
        print("Edit completed successfully!")
        return result, "✅ Image edited successfully!"
        
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        print(error_msg)
        import traceback
        traceback.print_exc()
        return None, error_msg

# Custom CSS
custom_css = """
.gradio-container {max-width: 1200px !important; margin: auto !important;}
#title {text-align: center; background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
-webkit-background-clip: text; -webkit-text-fill-color: transparent; font-size: 2.5em; font-weight: bold;}
"""

# Build interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    
    gr.HTML("<h1 id='title'>🎨 AI Image Editor</h1>")
    gr.Markdown("<p style='text-align: center; font-size: 1.2em;'>Edit images with simple text instructions</p>")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Upload Image", type="pil", height=400)
            
            edit_prompt = gr.Textbox(
                label="✏️ Edit Instructions",
                placeholder="Examples:\n- make it a cartoon\n- turn the sky to sunset\n- make it black and white\n- add snow on the ground",
                lines=3
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                num_steps = gr.Slider(10, 100, value=20, step=5, label="Steps (quality)")
                guidance_scale = gr.Slider(1, 10, value=7.5, step=0.5, label="Text Guidance")
                image_guidance_scale = gr.Slider(1, 2, value=1.5, step=0.1, label="Image Guidance")
            
            edit_btn = gr.Button("✨ Edit Image", variant="primary", size="lg")
        
        with gr.Column():
            output_image = gr.Image(label="Edited Image", type="pil", height=400)
            status = gr.Textbox(label="Status", interactive=False)
    
    gr.Markdown("""
    ### 💡 Example Prompts:
    - "make it look like a painting"
    - "turn day into night"
    - "make it winter"
    - "add a smile"
    - "make it black and white"
    - "turn it into a cartoon"
    """)
    
    gr.Markdown("---\n⚡ Powered by InstructPix2Pix")
    
    # Connect button
    edit_btn.click(
        fn=edit_image,
        inputs=[input_image, edit_prompt, num_steps, guidance_scale, image_guidance_scale],
        outputs=[output_image, status]
    )

if __name__ == "__main__":
    demo.queue(max_size=10)
    demo.launch()