image_generator.py

# from diffusers import StableDiffusionPipeline

# def generate_image(prompt):
#     model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
#     model.to("cuda")  # Use GPU for faster generation
#     image = model(prompt).images[0]
#     image.save("output.png")
#     return "output.png"

# if __name__ == "__main__":
#     prompt = "A friendly person saying 'How are you?'"
#     print("Generated Image Path:", generate_image(prompt))


# from diffusers import StableDiffusionPipeline
# import torch

# def generate_image(prompt):
#     model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
    
#     # Use GPU if available, otherwise fallback to CPU
#     device = "cuda" if torch.cuda.is_available() else "cpu"
#     model.to(device)
    
#     image = model(prompt).images[0]
#     image.save("output.png")
#     return "output.png"

# if __name__ == "__main__":
#     prompt = "A friendly person saying 'How are you?'"
#     print("Generated Image Path:", generate_image(prompt))

import spaces
from diffusers import StableDiffusionPipeline
import torch

# Preload the model globally
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1").to(device)

@spaces.GPU
def generate_image(prompt):
    """Generate an image based on the input prompt."""
    with torch.no_grad():
        image = pipeline(prompt).images[0]
    # Save the image locally and return the file path
    image_path = "generated_image.png"
    image.save(image_path)
    return image_path