mirror of https://github.com/skygpu/skynet.git
Drop old txt2img.py script
parent
cb92aed51c
commit
1fc2020ed5
|
@ -1,30 +0,0 @@
|
|||
import os
|
||||
import sys
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
|
||||
from huggingface_hub import login
|
||||
|
||||
assert torch.cuda.is_available()
|
||||
|
||||
login(token=os.environ['HF_TOKEN'])
|
||||
|
||||
pipe = StableDiffusionPipeline.from_pretrained(
|
||||
"runwayml/stable-diffusion-v1-5",
|
||||
torch_dtype=torch.float16,
|
||||
revision="fp16"
|
||||
)
|
||||
pipe = pipe.to("cuda")
|
||||
|
||||
prompt = sys.argv[1]
|
||||
image = pipe(
|
||||
prompt,
|
||||
width=640,
|
||||
height=640,
|
||||
guidance_scale=7.5, num_inference_steps=120
|
||||
).images[0]
|
||||
|
||||
image.save("/outputs/img.png")
|
Loading…
Reference in New Issue