File size: 972 Bytes
38b4eff | 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 32 33 | import torch
from diffusers import StableDiffusionPipeline
import os
import sys
def generate_image(prompt, output_path):
model_id = "runwayml/stable-diffusion-v1-5"
print(f"Loading model {model_id}...")
# Using float16 and CPU offload if possible, or just CPU if no CUDA
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float32 if device == "cpu" else torch.float16
)
pipe = pipe.to(device)
print(f"Generating image for prompt: {prompt}")
image = pipe(prompt, num_inference_steps=20).images[0]
image.save(output_path)
print(f"Image saved to {output_path}")
if __name__ == "__main__":
if len(sys.argv) < 3:
print("Usage: python neuralai_diffusion.py <prompt> <output_path>")
sys.exit(1)
prompt = sys.argv[1]
output_path = sys.argv[2]
generate_image(prompt, output_path)
|