import os
from diffusers import DiffusionPipeline
from PIL import Image

class EndpointHandler:
    def __init__(self, model_dir="stabilityai/stable-diffusion-x4-upscaler"):
        # Load the diffusion pipeline
        self.pipeline = DiffusionPipeline.from_pretrained(model_dir)

        # Ensure the output directory exists
        os.makedirs('output', exist_ok=True)

    def upscale_image(self, image_path):
        # Load the image
        image = Image.open(image_path)
        
        # Upscale the image
        upscaled_image = self.pipeline(image)
        
        # Save the upscaled image
        save_path = "output/upscaled_image.png"
        upscaled_image.save(save_path)
        
        return upscaled_image, save_path

# Example usage for testing
if __name__ == "__main__":
    handler = EndpointHandler()
    test_image_path = 'path_to_test_image.jpg'  # Replace with your test image path
    upscaled_image, save_path = handler.upscale_image(test_image_path)
    print(f"Upscaled image saved at {save_path}")