Spaces:
Running
Running
File size: 2,622 Bytes
33dda45 bb5a422 33dda45 bb5a422 33dda45 bb5a422 6642f4e bb5a422 33dda45 bb5a422 33dda45 6642f4e 33dda45 bb5a422 33dda45 |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch
class ImageGenerator:
def __init__(self, model_id="stabilityai/stable-diffusion-2-1-base", device="cuda"):
"""
Initialize the image generator with a specific model.
Args:
model_id (str): The model identifier for the stable diffusion model.
Default is "stabilityai/stable-diffusion-2-1-base".
device (str): The device to run the model on, either "cuda" or "cpu".
Default is "cuda".
"""
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
self.pipe = StableDiffusionPipeline.from_pretrained(
model_id,
scheduler=scheduler,
torch_dtype=torch.float16
)
self.pipe = self.pipe.to(device)
self.positive_prompt = "simple, icon"
self.negative_prompt = "3d, blurry, complex geometry, realistic"
def generate(self, prompt, negative_prompt=None, output_path=None, num_images=1, num_inference_steps=50):
"""
Generate an image based on the provided prompt.
Args:
prompt (str): The text description to generate an image from.
negative_prompt (str, optional): Elements to avoid in the generated image.
If None, uses the default negative prompt.
output_path (str, optional): Path to save the generated image.
If None, the image is not saved to disk.
num_images (int, optional): Number of images to generate.
Returns:
list[PIL.Image.Image]: The generated images.
"""
prompt = f"{prompt}, {self.positive_prompt}"
if negative_prompt is None:
negative_prompt = self.negative_prompt
images = self.pipe(
prompt,
negative_prompt=negative_prompt,
num_inference_steps=50,
num_images_per_prompt=num_images
).images
if output_path:
for i, image in enumerate(images):
image.save(f".cache/{output_path.replace('.png', f'_{i}.png')}")
return images
# Example usage
if __name__ == "__main__":
generator = ImageGenerator()
import time
start_time = time.time()
image = generator.generate(
prompt="magenta trapezoids layered on a transluscent silver sheet",
output_path="sheet.png",
num_images=4
)
end_time = time.time()
print(f"Time taken: {end_time - start_time} seconds")
|