GAN / app.py
akazmi's picture
Update app.py
da36322 verified
raw
history blame
787 Bytes
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(
"CompVis/stable-diffusion-v1-4",
use_auth_token=api_key
)
# Move the pipeline to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline.to(device)
# Reduce the image resolution for faster generation
pipeline.set_format("jpeg")
pipeline.set_resolution(512) # Reduce resolution (e.g., 512x512)
# Generate an image
prompt = "A serene lake surrounded by mountains during sunset"
image = pipeline(prompt).images[0]
# Save the generated image
image.save("generated_image.png")
print("Image generated and saved successfully!")