import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image

# Load model (will be cached in the Space, not your local GPU)
pipe = DiffusionPipeline.from_pretrained(
    "yahoo-inc/photo-background-generation", torch_dtype="float16"
).to("cuda")

def generate_background(image, mask, prompt):
    result = pipe(
        image=image,
        mask_image=mask,
        prompt=prompt,
        output_type="pil"
    ).images[0]
    return result

demo = gr.Interface(
    fn=generate_background,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Image(type="pil", label="Frame Mask (white = frame)"),
        gr.Textbox(label="Prompt")
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Salient Object-Aware Background Generation",
    description="Upload an image with a frame, its mask, and describe how the background should look."
)

demo.launch()