Spaces:
Sleeping
Sleeping
File size: 473 Bytes
4bce2e7 9c102f5 4bce2e7 db9ea67 4bce2e7 9c102f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | from diffusers import StableDiffusionPipeline
import torch
MODEL_NAME = "segmind/tiny-sd"
pipe = StableDiffusionPipeline.from_pretrained(
MODEL_NAME,
cache_dir="./models",
# local_files_only=True,
torch_dtype=torch.float32,
safety_checker=None,
)
pipe.enable_attention_slicing()
pipe.to("cpu")
@torch.inference_mode()
def generate_image(prompt):
image = pipe(prompt, height=512, width=512, num_inference_steps=10).images[0]
return image
|