import torch
from diffusers import StableDiffusionImg2ImgPipeline

import os
TOKEN = os.environ.get('TOKEN_SD') or True

torch.cuda.empty_cache()
device = "cpu"


def pipe_image(prompt,
               init_image=None,
               strength=0.6):
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", 
                                                revision="fp16", 
                                                torch_dtype=torch.float32,
                                                use_auth_token=TOKEN
                                                )
    pipe = pipe.to(device)
    #pipe.enable_attention_slicing()
    images = pipe(prompt=prompt,
                    init_image=init_image,
                    strength=strength,
                    guidance_scale=7.5,
                    )
    return images