Scriptyfa / utils.py
Ridam7777u's picture
Upload 5 files
be416b0 verified
from diffusers import DiffusionPipeline
import torch
prompt_presets = {
"Cinematic": "cinematic lighting, epic composition, 8k",
"Realistic Portrait": "realistic face, shallow depth of field, photography",
"Anime Style": "anime, cel-shading, crisp lines, colorful",
"Fantasy": "mythical, magical light, detailed, fantasy world",
"None": ""
}
def load_model():
print("Loading SDXL model...")
pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
variant="fp16" if torch.cuda.is_available() else None
)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")
return pipe
def generate_image(pipe, prompt: str, guidance: float, steps: int, width: int, height: int):
if not prompt or len(prompt.strip()) < 5:
raise ValueError("Prompt too short. Please describe your idea better.")
result = pipe(prompt, guidance_scale=guidance, num_inference_steps=steps, height=height, width=width)
return result.images[0]