Update app.py

app.py

@@ -1,3 +1,5 @@
+import argparse
+import os
 import random
 import torch
 import numpy as np
@@ -6,21 +8,13 @@ from PIL import Image
 import re
 
 def generate_image(pipe, prompt, seed=42, randomize_seed=True, width=768, height=768, guidance_scale=4.5, num_inference_steps=20):
-    """
-    使用 FLUX.1-Krea-dev 模型生成图像。
-
-    Args:
-        pipe: 配置好的 Diffusers pipeline.
-
-        prompt (str): 文本提示.
-        seed (int): 随机种子.
-        randomize_seed (bool): 是否随机化种子.
-    generator = torch.Generator(device=pipe.device).manual_seed(seed)
+    MAX_SEED = np.iinfo(np.int32).max
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+        generator = torch.Generator(device=pipe.device).manual_seed(seed)
+        print(f"ℹ️  使用种子: {seed}")
     
-    print(f"ℹ️  使用种子: {seed}")
     print("🚀 开始生成图像...")
-    
-    # 直接调用 pipeline 生成 PIL 图像，内部会自动处理解码
     image = pipe(
         prompt=prompt,
         guidance_scale=guidance_scale,
@@ -29,4 +23,61 @@ def generate_image(pipe, prompt, seed=42, randomize_seed=True, width=768, height
         height=height,
         generator=generator,
         output_type="pil"
-    ).images[0]
+    ).images[0]
+    
+    return image, seed
+
+def main():
+    parser = argparse.ArgumentParser(description="使用 FLUX.1-Krea-dev 模型从文本提示生成图像。")
+    parser.add_argument("--prompt", type=str, required=True, help="用于图像生成的文本提示。")
+    parser.add_argument("--seed", type=int, default=None, help="随机种子。如果未提供，将随机生成。")
+    parser.add_argument("--steps", type=int, default=20, help="推理步数。")
+    parser.add_argument("--width", type=int, default=768, help="图像宽度。")
+    parser.add_argument("--height", type=int, default=768, help="图像高度。")
+    parser.add_argument("--guidance", type=float, default=4.5, help="指导比例 (Guidance Scale)。")
+    args = parser.parse_args()
+
+    print("⏳ 正在加载模型，请稍候...")
+    dtype = torch.bfloat16
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", subfolder="vae", torch_dtype=dtype)
+    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", torch_dtype=dtype, vae=good_vae).to(device)
+
+    if device == "cuda":
+        torch.cuda.empty_cache()
+
+    print(f"✅ 模型加载完成，使用设备: {device}")
+
+    print(f"🎨 开始为提示生成图像: '{args.prompt}'")
+
+    randomize = args.seed is None
+    seed_value = args.seed if not randomize else 42 
+
+    generated_image, used_seed = generate_image(
+        pipe=pipe,
+        prompt=args.prompt,
+        seed=seed_value,
+        randomize_seed=randomize,
+        width=args.width,
+        height=args.height,
+        num_inference_steps=args.steps,
+        guidance_scale=args.guidance
+    )
+
+    output_dir = "output"
+    os.makedirs(output_dir, exist_ok=True)
+    
+    safe_prompt = re.sub(r'[^\w\s-]', '', args.prompt).strip()
+    safe_prompt = re.sub(r'[-\s]+', '_', safe_prompt)
+    
+    filename = f"{safe_prompt[:50]}_{used_seed}.png"
+    filepath = os.path.join(output_dir, filename)
+
+    print(f"💾 正在保存图像到: {filepath}")
+    generated_image.save(filepath)
+
+    print("🎉 完成！")
+
+if __name__ == "__main__":
+    main()