inference.py

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline, DDIMScheduler
import sys


if len(sys.argv) > 1:
    model_path = sys.argv[1]
else:
    model_path = 'fine-tuned-model-output/800'

if len(sys.argv) > 2:
    prompt = sys.argv[2]
else:
    prompt = "a photo of sks dog"

if len(sys.argv) > 3:
    outputFileName = sys.argv[3]
else:
    outputFileName = "output.png"
print("prompot: ", prompt)
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
#@markdown Can set random seed here for reproducibility.
g_cuda = torch.Generator(device='cuda')
seed = 52362 #@param {type:"number"}
g_cuda.manual_seed(seed)

 #@param {type:"string"}
negative_prompt = "" #@param {type:"string"}
num_samples = 1 #@param {type:"number"}
guidance_scale = 7.5 #@param {type:"number"}
num_inference_steps = 60 #@param {type:"number"}
height = 512 #@param {type:"number"}
width = 512 #@param {type:"number"}

with autocast("cuda"), torch.inference_mode():
    images = pipe(
        prompt,
        height=height,
        width=width,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_samples,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=g_cuda
    ).images
i = 0
for img in images:
    # save image to disk
    img.save(outputFileName)
    i+=1