Ai image generator

 

from diffusers import StableDiffusionPipeline
import torch

def generate_image(prompt, save_path="generated_image.png"):
    model_id = "runwayml/stable-diffusion-v1-5"
    pipe = StableDiffusionPipeline.from_pretrained(model_id)
    pipe.to("cuda" if torch.cuda.is_available() else "cpu")

    image = pipe(prompt).images[0]
    image.save(save_path)
    print(f"Image saved to {save_path}")

if __name__ == "__main__":
    user_prompt = input("Enter your image description: ")
    generate_image(user_prompt)