import spaces
import gradio as gr
from rembg import remove
from PIL import Image
import torch
from diffusers import StableDiffusionPipeline

# Load the Stable Diffusion pipeline (requires GPU)
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
).to("cuda")

@spaces.GPU  # Decorate the function to enable ZeroGPU dynamic GPU allocation
def replace_background(profile_img, prompt):
    if profile_img is None or not prompt.strip():
        return "Please upload a profile picture and enter a background prompt."

    # Step 1: Remove background (keeps person only)
    profile_img = profile_img.convert("RGBA")
    person_only = remove(profile_img).resize((512, 512))

    # Step 2: Generate new background image from prompt
    background = pipe(prompt, height=512, width=512).images[0].convert("RGBA")

    # Step 3: Composite person onto new background
    final_image = Image.alpha_composite(background, person_only)
    return final_image

with gr.Blocks() as demo:
    gr.Markdown("# 🧍 AI Background Replacement App")
    gr.Markdown(
        "Upload a profile picture, enter a prompt like 'tropical island at sunset', and get a version with a new AI-generated background!"
    )

    with gr.Row():
        image_input = gr.Image(label="Upload Profile Image", type="pil")
        prompt_input = gr.Textbox(label="Background Prompt", placeholder="e.g., a starry night in Paris")

    output_image = gr.Image(label="Result", type="pil")
    generate_button = gr.Button("Generate New Background")

    generate_button.click(fn=replace_background, inputs=[image_input, prompt_input], outputs=output_image)

    gr.Markdown("⚠️ Works best with a portrait image and short background prompts.\n🚀 GPU required for best performance.")

if __name__ == "__main__":
    demo.launch()