File size: 663 Bytes
4398c5e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
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