import torch
from diffusers import StableDiffusionPipeline

# Hugging Face API key
api_key = "your_huggingface_api_key"

# Load the Stable Diffusion pipeline with the API key
pipeline = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2",  # You can try using the newer and faster version of Stable Diffusion.
    use_auth_token=api_key
)

# Move the pipeline to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline.to(device)

# Define a smaller resolution for faster generation
height = 512  # You can try 512 or 256 to make it faster
width = 512

# Generate an image
prompt = "A serene lake surrounded by mountains during sunset"
image = pipeline(prompt, height=height, width=width).images[0]

# Save the generated image
image.save("generated_image.png")
print("Image generated and saved successfully!")