from diffusers import DiffusionPipeline


class SDXL:

    def __init__(self):
        pass

    def load_model(self):
        pipeline = DiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V6.0_B1_noVAE",)
        pipeline.safety_checker = None

        return pipeline

    def predict(self, pipeline, prompt):
        images = pipeline(prompt=prompt,).images[0]
        
        return images