Update app.py

app.py

@@ -3,17 +3,16 @@ import os
 import random
 import torch
 import numpy as np
-from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
+from diffusers import DiffusionPipeline, AutoencoderKL
 from PIL import Image
 import re
 
-def generate_image(pipe, good_vae, prompt, seed=42, randomize_seed=True, width=768, height=768, guidance_scale=4.5, num_inference_steps=20):
+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.
-        good_vae: 高质量的 VAE 解码器.
+        pipe: 配置好的 Diffusers pipeline.
         prompt (str): 文本提示.
         seed (int): 随机种子.
         randomize_seed (bool): 是否随机化种子.
@@ -32,30 +31,18 @@ def generate_image(pipe, good_vae, prompt, seed=42, randomize_seed=True, width=7
     generator = torch.Generator(device=pipe.device).manual_seed(seed)
     
     print(f"ℹ️  使用种子: {seed}")
-    print("1. 正在生成潜在向量 (latents)...")
+    print("🚀 开始生成图像...")
     
-    # 使用 pipeline 生成潜在向量
-    latents = pipe(
+    # 直接调用 pipeline 生成 PIL 图像，内部会自动处理解码
+    image = pipe(
         prompt=prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         width=width,
         height=height,
         generator=generator,
-        output_type="latent"
-    ).images
-
-    print("2. 使用高质量 VAE 解码图像...")
-    
-    # 使用高质量的 VAE 解码潜在向量
-    # 需要根据 VAE 的配置进行缩放
-    latents = latents / good_vae.config.scaling_factor
-    image_tensor = good_vae.decode(latents, return_dict=False)[0]
-
-    print("3. 后处理图像...")
-    
-    # 将张量转换为 PIL 图像
-    image = pipe.image_processor.postprocess(image_tensor, output_type="pil")[0]
+        output_type="pil"
+    ).images[0]
     
     return image, seed
 
@@ -70,6 +57,7 @@ def main():
     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()
 
     # --- 模型加载 ---
@@ -77,12 +65,11 @@ def main():
     dtype = torch.bfloat16
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
-    # 加载两个 VAE：一个用于快速预览（在 pipeline 中），一个用于高质量最终输出
-    taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
-    good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", subfolder="vae", torch_dtype=dtype).to(device)
+    # 加载高质量的 VAE 解码器
+    good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", subfolder="vae", torch_dtype=dtype)
     
-    # 加载主 pipeline，并指定使用较小的 VAE 进行快速潜在向量生成
-    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", torch_dtype=dtype, vae=taef1).to(device)
+    # 加载主 pipeline，并直接将高质量的 VAE 传入
+    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()
@@ -90,30 +77,32 @@ def main():
     print(f"✅ 模型加载完成，使用设备: {device}")
 
     # --- 图像生成 ---
-    print(f"🚀 开始为提示生成图像: '{args.prompt}'")
+    print(f"🎨 开始为提示生成图像: '{args.prompt}'")
     
     randomize = args.seed is None
-    seed_value = args.seed if not randomize else 42 # 如果指定了种子则使用，否则 generate_image 会随机生成
+    # 如果用户没有指定种子，则在调用函数时随机化；否则使用用户指定的种子
+    seed_value = args.seed if not randomize else 42 
 
     generated_image, used_seed = generate_image(
         pipe=pipe,
-        good_vae=good_vae,
         prompt=args.prompt,
         seed=seed_value,
         randomize_seed=randomize,
         width=args.width,
         height=args.height,
-        num_inference_steps=args.steps
+        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)