| |
| !pip install torch torchvision transformers diffusers |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from diffusers import StableDiffusionPipeline |
| import matplotlib.pyplot as plt |
|
|
| |
| model_id = "CompVis/stable-diffusion-v1-4" |
| pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16) |
| pipe = pipe.to("cuda") |
|
|
| |
| def generate_animal_image(prompt, output_path="animal_image.png"): |
| with torch.no_grad(): |
| image = pipe(prompt).images[0] |
| image.save(output_path) |
| plt.imshow(image) |
| plt.axis('off') |
| plt.show() |
| print(f"Generated image saved as '{output_path}'") |
|
|
| |
| prompt = "a cute animal in a forest" |
| generate_animal_image(prompt) |
|
|