Code Editor for app.py

import streamlit as st
from PIL import Image
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
import io

st.set_page_config(page_title="Face Swap AI", page_icon="👤")

st.title("Face Swap AI App")
st.write("Upload two images to swap faces using AI")

# Initialize the model
@st.cache_resource
def load_model():
    model_id = "timbrooks/instruct-pix2pix"
    pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
        safety_checker=None
    )
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    if torch.cuda.is_available():
        pipe = pipe.to("cuda")
    return pipe

try:
    pipe = load_model()
    st.success("✅ Model loaded successfully!")
except Exception as e:
    st.error(f"Error loading model: {str(e)}")

# Image upload
col1, col2 = st.columns(2)

with col1:
    source_image = st.file_uploader("Upload source image (face to be replaced)", type=["jpg", "jpeg", "png"])
    if source_image:
        st.image(source_image, caption="Source Image", use_column_width=True)

with col2:
    target_image = st.file_uploader("Upload target image (face to be used)", type=["jpg", "jpeg", "png"])
    if target_image:
        st.image(target_image, caption="Target Image", use_column_width=True)

def process_image(image):
    if image is not None:
        img = Image.open(image)
        # Resize image to 512x512 (model requirement)
        img = img.resize((512, 512))
        return img
    return None

def swap_faces():
    if source_image is None or target_image is None:
        st.warning("Please upload both images first!")
        return
    
    with st.spinner("Processing... This may take a while"):
        try:
            # Process images
            source = process_image(source_image)
            target = process_image(target_image)
            
            # Generate face swap instruction
            instruction = "Replace the face in the image with the face from the reference image"
            
            # Generate image
            output = pipe(
                instruction,
                image=source,
                num_inference_steps=20,
                image_guidance_scale=1
            ).images[0]
            
            # Display result
            st.success("Face swap completed!")
            st.image(output, caption="Result", use_column_width=True)
            
            # Add download button
            buf = io.BytesIO()
            output.save(buf, format="PNG")
            btn = st.download_button(
                label="Download result",
                data=buf.getvalue(),
                file_name="face_swap_result.png",
                mime="image/png"
            )
            
        except Exception as e:
            st.error(f"Error during face swap: {str(e)}")

if source_image and target_image:
    if st.button("Swap Faces"):
        swap_faces()

st.markdown("---")
st.markdown("""
### How to use:
1. Upload a source image (the image containing the face you want to replace)
2. Upload a target image (the image containing the face you want to use)
3. Click 'Swap Faces' to generate the result
4. Download the result using the download button

### Notes:
- Images will be resized to 512x512 pixels
- Processing may take a few minutes depending on your system
- For best results, use clear frontal face images
""")

# Add session state to store the results
if 'history' not in st.session_state:
    st.session_state.history = []