models/realistic_vision_v6b1.py

from diffusers import DiffusionPipeline

class RealV6B1:
    
    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