import torch from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler from PIL import Image import argparse import os class AIImageGenerator: def __init__(self): self.model_name = "stabilityai/stable-diffusion-2-1" self.device = "cuda" if torch.cuda.is_available() else "cpu" self.pipe = None self.load_model() def load_model(self): """Load the Stable Diffusion model with optimized settings""" print(f"Loading {self.model_name}...") self.pipe = StableDiffusionPipeline.from_pretrained( self.model_name, torch_dtype=torch.float16, safety_checker=None, use_safetensors=True ) self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config) self.pipe = self.pipe.to(self.device) self.pipe.enable_attention_slicing() # Reduces VRAM usage print("Model loaded successfully!") def generate_image(self, prompt, aspect_ratio="square", negative_prompt=None, output_path="output.png"): """Generate image based on parameters""" # Validate output path if not output_path.lower().endswith(('.png', '.jpg', '.jpeg')): output_path += ".png" # Set dimensions based on aspect ratio dimensions = { "square": (768, 768), "portrait": (640, 960), # 2:3 ratio "portrait_9_16": (540, 960), # 9:16 mobile ratio "landscape": (960, 640), # 3:2 ratio "landscape_16_9": (960, 540) # 16:9 widescreen }.get(aspect_ratio, (768, 768)) # Default negative prompt if not provided if negative_prompt is None: negative_prompt = ("blurry, low quality, distorted, bad anatomy, " "extra fingers, mutated hands, poorly drawn face") print(f"Generating {aspect_ratio} image ({dimensions[0]}x{dimensions[1]})...") # Generate the image image = self.pipe( prompt=prompt, negative_prompt=negative_prompt, width=dimensions[0], height=dimensions[1], num_inference_steps=50, guidance_scale=7.5, generator=torch.Generator(self.device).manual_seed(torch.seed()) ).images[0] # Save the image image.save(output_path) print(f"Image successfully saved to {os.path.abspath(output_path)}") return image def main(): parser = argparse.ArgumentParser(description="AI Image Generator with Multiple Aspect Ratios") parser.add_argument("--prompt", type=str, required=True, help="Text prompt for image generation") parser.add_argument("--aspect", type=str, default="square", choices=["square", "portrait", "portrait_9_16", "landscape", "landscape_16_9"], help="Aspect ratio of the image") parser.add_argument("--negative", type=str, default=None, help="Negative prompt (things to exclude from image)") parser.add_argument("--output", type=str, default="ai_generated_image.png", help="Output file path (supports .png, .jpg)") args = parser.parse_args() # Initialize generator generator = AIImageGenerator() # Generate and save image generated_image = generator.generate_image( prompt=args.prompt, aspect_ratio=args.aspect, negative_prompt=args.negative, output_path=args.output ) # Show the image generated_image.show() if __name__ == "__main__": main()