ImageCreationByAnees / image_mode.py
aneesqumar's picture
Create image_mode.py
0a169f3 verified
Raw
History Blame Contribute Delete
560 Bytes
import torch
from diffusers import StableDiffusionPipeline
# Load model once (important for performance)
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)
if torch.cuda.is_available():
pipe = pipe.to("cuda")
def generate_image(prompt: str):
if not prompt.strip():
return None
image = pipe(
prompt,
guidance_scale=7.5,
num_inference_steps=30
).images[0]
return image