File size: 403 Bytes
463297f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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