Update run_civitai_sdxl.py

run_civitai_sdxl.py

@@ -1,9 +1,13 @@
+#### inference code drived from https://civitai.com/models/833294 and https://huggingface.co/spaces/nyanko7/toaru-xl-model
+
 import torch
-from diffusers import StableDiffusionXLPipeline
-from diffusers import EulerDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
+from PIL import Image
+import random
 
+# 初始化 pipe，只执行一次
+ckpt_path = "nyaflow-xl-alpha.safetensors"  # https://huggingface.co/nyanko7/nyaflow-xl-alpha
 ckpt_path = "noobaiXLNAIXL_vPred10Version.safetensors" #### https://civitai.com/models/833294
-ckpt_path = "nyaflow-xl-alpha.safetensors" #### https://huggingface.co/nyanko7/nyaflow-xl-alpha
 pipe = StableDiffusionXLPipeline.from_single_file(
     ckpt_path,
     use_safetensors=True,
@@ -14,17 +18,56 @@ pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, **sch
 pipe.enable_xformers_memory_efficient_attention()
 pipe = pipe.to("cuda")
 
-prompt = """masterpiece, best quality,artist:john_kafka,artist:nixeu,artist:quasarcake, chromatic aberration, film grain, horror \(theme\), limited palette, x-shaped pupils, high contrast, color contrast, cold colors, arlecchino \(genshin impact\), black theme,  gritty, graphite \(medium\)"""
-prompt = "zhongli"
-negative_prompt = "nsfw, worst quality, old, early, low quality, lowres, signature, username, logo, bad hands, mutated hands, mammal, anthro, furry, ambiguous form, feral, semi-anthro"
-
-image = pipe(
-    prompt=prompt,
-    negative_prompt=negative_prompt,
-    width=832,
-    height=1216,
-    num_inference_steps=28,
-    guidance_scale=5,
-    generator=torch.Generator().manual_seed(42),
-).images[0]
-image
+# 定义默认参数
+PRESET_Q = "year_2022, best quality, high quality, very aesthetic"
+NEGATIVE_PROMPT = "lowres, worst quality, displeasing, bad anatomy, text, error, extra digit, cropped, error, fewer, extra, missing, worst quality, jpeg artifacts, censored, ai-generated worst quality displeasing, bad quality"
+
+def generate_image(
+    prompt: str,
+    preset: str = PRESET_Q,
+    height: int = 1216,
+    width: int = 832,
+    negative_prompt: str = NEGATIVE_PROMPT,
+    guidance_scale: float = 4.0,
+    randomize_seed: bool = True,
+    seed: int = 42,
+    inference_steps: int = 25,
+) -> Image:
+    # 合并 prompt 和 preset
+    prompt = prompt.strip() + ", " + preset.strip()
+    negative_prompt = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None
+
+    # 随机化种子
+    if randomize_seed:
+        seed = random.randint(0, 9007199254740991)
+
+    # 设置生成器
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    # 限制推理步数
+    if inference_steps > 50:
+        inference_steps = 50
+
+    # 生成图像
+    image = pipe(
+        prompt,
+        height=height,
+        width=width,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        generator=generator,
+        num_inference_steps=inference_steps
+    ).images[0]
+
+    return image
+
+
+# 示例调用
+if __name__ == "__main__":
+    prompt = "zhongli"
+    image = generate_image(prompt)
+    image
+
+    prompt = "Neuvillette"
+    image = generate_image(prompt)
+    image