from stable_diffusion_videos import StableDiffusionWalkPipeline, Interface
from diffusers.schedulers import LMSDiscreteScheduler
import torch

import os
from huggingface_hub import HfFolder

HfFolder().save_token(os.environ['HF_TOKEN'])

pipeline = StableDiffusionWalkPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    torch_dtype=torch.float16,
    revision="fp16",
    scheduler=LMSDiscreteScheduler(
        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
    ),
    safety_checker=None
).to("cuda")
pipeline.enable_attention_slicing()

interface = Interface(pipeline)

if __name__ == '__main__':
    interface.launch()