Spaces:
Sleeping
Sleeping
| """ | |
| AI Image Generation Module | |
| Uses Stable Diffusion to generate images from text prompts | |
| """ | |
| import torch | |
| from diffusers import StableDiffusionPipeline | |
| from PIL import Image | |
| import os | |
| class ImageGenerator: | |
| def __init__(self, model_id="stabilityai/stable-diffusion-2-1-base"): | |
| """Initialize the Stable Diffusion pipeline""" | |
| self.model_id = model_id | |
| self.pipe = None | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| def load_model(self): | |
| """Load the Stable Diffusion model""" | |
| if self.pipe is None: | |
| print(f"Loading model on {self.device}...") | |
| self.pipe = StableDiffusionPipeline.from_pretrained( | |
| self.model_id, | |
| torch_dtype=torch.float16 if self.device == "cuda" else torch.float32, | |
| safety_checker=None | |
| ) | |
| self.pipe = self.pipe.to(self.device) | |
| # Enable memory optimizations | |
| if self.device == "cuda": | |
| self.pipe.enable_attention_slicing() | |
| print("Model loaded successfully!") | |
| def generate( | |
| self, | |
| prompt: str, | |
| negative_prompt: str = "", | |
| num_inference_steps: int = 30, | |
| guidance_scale: float = 7.5, | |
| width: int = 512, | |
| height: int = 512, | |
| seed: int = None | |
| ) -> Image.Image: | |
| """ | |
| Generate an image from a text prompt | |
| Args: | |
| prompt: Text description of desired image | |
| negative_prompt: What to avoid in the image | |
| num_inference_steps: Number of denoising steps (higher = better quality, slower) | |
| guidance_scale: How closely to follow the prompt (7-10 recommended) | |
| width: Image width (must be multiple of 8) | |
| height: Image height (must be multiple of 8) | |
| seed: Random seed for reproducibility | |
| Returns: | |
| PIL Image | |
| """ | |
| self.load_model() | |
| # Set seed for reproducibility | |
| generator = None | |
| if seed is not None: | |
| generator = torch.Generator(device=self.device).manual_seed(seed) | |
| # Generate image | |
| with torch.inference_mode(): | |
| result = self.pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale, | |
| width=width, | |
| height=height, | |
| generator=generator | |
| ) | |
| return result.images[0] | |
| def unload_model(self): | |
| """Free up memory by unloading the model""" | |
| if self.pipe is not None: | |
| del self.pipe | |
| self.pipe = None | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| print("Model unloaded") | |
| # Example usage | |
| if __name__ == "__main__": | |
| generator = ImageGenerator() | |
| image = generator.generate( | |
| prompt="A fantasy landscape with mountains and a castle at sunset", | |
| seed=42 | |
| ) | |
| image.save("test_generated.png") | |
| print("Image saved as test_generated.png") | |