app.py

import streamlit as st
from diffusers import StableDiffusionInpaintPipeline, StableDiffusionPipeline
from PIL import Image
import torch
import uuid
import os

# ---------------------------------------
# Load Models (Works on HF Spaces CPU)
# ---------------------------------------

@st.cache_resource
def load_inpaint_model():
    return StableDiffusionInpaintPipeline.from_pretrained(
        "runwayml/stable-diffusion-inpainting",
        torch_dtype=torch.float32
    )

@st.cache_resource
def load_txt2img_model():
    return StableDiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-1",
        torch_dtype=torch.float32
    )

# Load both
inpaint_model = load_inpaint_model()
txt2img_model = load_txt2img_model()


# ---------------------------------------
# Streamlit UI
# ---------------------------------------

st.set_page_config(
    page_title="AI Image Editor",
    page_icon="🎨",
    layout="wide"
)

st.title("🎨 AI Image Editor (HuggingFace + Streamlit)")
st.write("Upload an image and edit it using natural language prompts.")

# Sidebar — History Manager
st.sidebar.header("🕑 Edit History")
if "history" not in st.session_state:
    st.session_state.history = []

if "current_image" not in st.session_state:
    st.session_state.current_image = None


# ---------------------------------------
# Image Upload
# ---------------------------------------

uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

if uploaded_file:
    input_image = Image.open(uploaded_file).convert("RGB")
    st.session_state.current_image = input_image
    st.image(input_image, caption="Original Image", use_column_width=True)

# ---------------------------------------
# Prompt Input
# ---------------------------------------

prompt = st.text_input("Enter editing instruction (e.g., 'remove background', 'add sunset sky', 'cartoon style')")

edit_button = st.button("Apply Edit")


# ---------------------------------------
# Editing Logic
# ---------------------------------------

def save_to_history(image, prompt):
    st.session_state.history.append({
        "prompt": prompt,
        "image": image
    })

def display_history():
    for i, step in enumerate(st.session_state.history):
        st.sidebar.image(step["image"], caption=f"Step {i+1}: {step['prompt']}", width=150)


if edit_button and st.session_state.current_image is not None:

    with st.spinner("Processing..."):

        # Decide editing type
        if "remove background" in prompt.lower():
            # Replace background with white (simple trick)
            edited_image = st.session_state.current_image.convert("RGBA")
            datas = edited_image.getdata()
            new_data = []
            for item in datas:
                if item[0] > 200 and item[1] > 200 and item[2] > 200:
                    new_data.append((255, 255, 255, 0))
                else:
                    new_data.append(item)
            edited_image.putdata(new_data)

        elif "add" in prompt.lower() or "change" in prompt.lower():
            # Use text-to-image model for creative edits
            edited_image = txt2img_model(prompt=prompt).images[0]

        else:
            # Default = inpainting logic
            mask = Image.new("L", st.session_state.current_image.size, 0)
            edited_image = inpaint_model(
                prompt=prompt,
                image=st.session_state.current_image,
                mask_image=mask
            ).images[0]

        st.session_state.current_image = edited_image
        save_to_history(edited_image, prompt)

        st.image(edited_image, caption="Edited Image", use_column_width=True)


# ---------------------------------------
# Show History
# ---------------------------------------

display_history()

# ---------------------------------------
# Download Button
# ---------------------------------------

if st.session_state.current_image:
    img_name = f"edited_{uuid.uuid4().hex}.png"
    st.download_button(
        label="Download Edited Image",
        data=st.session_state.current_image.save(img_name),
        file_name=img_name,
        mime="image/png"
    )