import gradio as gr
from diffusers import DiffusionPipeline
import torch

# Initialize the pipeline
pipe = DiffusionPipeline.from_pretrained("Lookingsoft-team/object_detection")
if torch.cuda.is_available():
    pipe = pipe.to("cuda")


def detect_objects(image):
    if image is None:
        return None

    # Process the image through the pipeline
    # Note: This is a placeholder - actual processing will depend on the model's specific requirements
    results = pipe(image=image)

    # Return the processed image with detections
    return results.images[0]


# Create Gradio interface
demo = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Object Detection",
    description="Upload an image and the model will detect objects in it!",
)

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