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)