man08man's picture
Upload 3 files
4398c5e verified
raw
history blame
663 Bytes
from diffusers import StableDiffusionPipeline
import torch
import random
def generate_image(prompt):
# Use updated model loading syntax
pipe = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1",
torch_dtype=torch.float16, # Add this for better performance
use_auth_token="YOUR_HF_TOKEN", # Replace with your token
safety_checker=None # Optional: disable safety filter
)
# Generate random seed
seed = random.randint(0, 1000000)
generator = torch.Generator().manual_seed(seed)
# Generate image
image = pipe(prompt, generator=generator).images[0]
return image