Skip to content

Instantly share code, notes, and snippets.

@maxim04

maxim04/inference.py

Last active February 1, 2023 20:37
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save maxim04/320c7adec2713b32aba9a17f84f6c2eb to your computer and use it in GitHub Desktop.
Save maxim04/320c7adec2713b32aba9a17f84f6c2eb to your computer and use it in GitHub Desktop.
Download ZIP
Raw
inference.py
import os
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline, DDIMScheduler
model_path = os.getenv('WEIGHTS_DIR')
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
pipe = StableDiffusionPipeline.from_pretrained(model_path, scheduler=scheduler, safety_checker=None, torch_dtype=torch.float16).to("cuda")
g_cuda = None
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment