File size: 1,027 Bytes
2c57eba |
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 26 27 28 29 30 31 |
from diffusers import StableDiffusionPipeline # Import the StableDiffusionPipeline class from diffusers
import torch
# Define the model ID for Stable Diffusion
model_id = "CompVis/stable-diffusion-v1-4"
# Load the pre-trained model
pipe = StableDiffusionPipeline.from_pretrained(model_id)
# Move the model to GPU if available, otherwise use CPU
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
# Function to generate images based on text prompts
def generate_image(prompt, output_path="generated_image.png"):
# Generate the image
image = pipe(prompt).images[0]
# Save the generated image
image.save(output_path)
print(f"Image saved at: {output_path}")
# Example prompts for image generation
prompt1 = "A futuristic city skyline at sunset"
generate_image(prompt1, "image1.png")
prompt2 = "A beautiful landscape from mountains during sunrise"
generate_image(prompt2, "image2.png")
prompt3 = "A beautiful landscape from mountains during sunset"
generate_image(prompt3, "image3.png")
|