Update app.py

app.py

@@ -1,154 +1,394 @@
 import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
 import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import datetime
+import io
+import json
+import os
+from typing import Optional
+
+# ======================
+# 配置区（你只需修改这里即可扩展）
+# ======================
+
+# 1. 基础模型
+BASE_MODEL = "SG161222/RealisticVisionV6.0"
+
+# 2. 固定LoRA（不可选，自动加载）
+FIXED_LORAS = [
+    ("Lykon/epiCRealism_LoRA", 0.8),      # 质量增强
+    ("latent-consistency/lora-dreamshaper", 0.7),  # 姿势控制
+]
+
+# 3. 风格模板（自动拼接到用户提示词前）
+STYLE_PROMPTS = {
+    "None": "",
+    "Realistic": "photorealistic, ultra-detailed skin, natural lighting, 8k, professional photography, f/1.8, shallow depth of field, Canon EOS R5, ",
+    "Anime": "anime style, cel shading, vibrant colors, detailed eyes, studio ghibli, trending on pixiv, ",
+    "Comic": "comic book style, bold outlines, dynamic angles, comic panel, Marvel style, inked lines, ",
+    "Watercolor": "watercolor painting, soft brush strokes, translucent layers, artistic, painterly, paper texture, ",
+}
 
+# 4. 可选LoRA下拉菜单（用户可选1个，None表示清除）
+OPTIONAL_LORAS = [
+    "None",
+    "Add Detail: https://huggingface.co/latent-consistency/lora-add-detail",
+    "Vintage Photo: https://huggingface.co/ckpt/LoRA-vintage-photo",
+    "Cinematic: https://huggingface.co/latent-consistency/lora-cinematic",
+    "Portrait Enhancer: https://huggingface.co/deforum/Portrait-Enhancer-LoRA",
+    "Soft Focus: https://huggingface.co/latent-consistency/lora-soft-focus",
+]
+
+# 解析可选LoRA的名称和ID
+OPTIONAL_LORA_MAP = {}
+for item in OPTIONAL_LORAS:
+    if item != "None":
+        name, url = item.split(": ", 1)
+        OPTIONAL_LORA_MAP[name] = url
+    else:
+        OPTIONAL_LORA_MAP["None"] = None
+
+# 默认参数
+DEFAULT_SEED = -1
+DEFAULT_WIDTH = 1024
+DEFAULT_HEIGHT = 1024
+DEFAULT_LORA_SCALE = 0.8
+DEFAULT_STEPS = 30
+DEFAULT_CFG = 7.5
+
+# ======================
+# 全局变量：延迟加载模型
+# ======================
+pipe = None
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
+
+def load_pipeline():
+    global pipe
+    if pipe is None:
+        print("🚀 Loading base model...")
+        pipe = StableDiffusionPipeline.from_pretrained(
+            BASE_MODEL,
+            torch_dtype=torch.float16,
+            safety_checker=None,
+            requires_safety_checker=False,
+        ).to(device)
+        pipe.enable_attention_slicing()
+        pipe.enable_vae_slicing()
+        pipe.enable_model_cpu_offload()  # 适配ZeroGPU
+        print("✅ Base model loaded.")
+    return pipe
+
+def unload_pipeline():
+    global pipe
+    if pipe is not None:
+        del pipe
+        torch.cuda.empty_cache()
+        pipe = None
+        print("🗑️ Pipeline unloaded.")
+
+# ======================
+# 主生成函数
+# ======================
+def generate_image(
+    prompt, negative_prompt, style, seed, width, height, optional_lora_name, lora_scale,
+    steps, cfg_scale
 ):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    global pipe
+
+    # 加载模型（懒加载）
+    pipe = load_pipeline()
+
+    # 处理种子
+    if seed == -1:
+        seed = torch.randint(0, 2**32, (1,)).item()
+    generator = torch.Generator(device=device).manual_seed(seed)
 
-    generator = torch.Generator().manual_seed(seed)
+    # 拼接风格提示词
+    full_prompt = STYLE_PROMPTS[style] + prompt
+    full_negative_prompt = negative_prompt
 
+    # 加载固定LoRA（每次生成前都加载，确保状态正确）
+    for lora_id, scale in FIXED_LORAS:
+        pipe.load_lora_weights(lora_id, adapter_name=lora_id)
+        pipe.set_adapters([lora_id], adapter_weights=[scale])
+
+    # 加载可选LoRA（如果非None）
+    if optional_lora_name != "None":
+        lora_url = OPTIONAL_LORA_MAP[optional_lora_name]
+        pipe.load_lora_weights(lora_url, adapter_name=optional_lora_name)
+        pipe.set_adapters([lora_id for lora_id, _ in FIXED_LORAS] + [optional_lora_name], 
+                          adapter_weights=[scale for _, scale in FIXED_LORAS] + [lora_scale])
+    else:
+        # 清除所有可选LoRA，只保留固定
+        pipe.set_adapters([lora_id for lora_id, _ in FIXED_LORAS], 
+                          adapter_weights=[scale for _, scale in FIXED_LORAS])
+
+    # 生成图像
     image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
+        prompt=full_prompt,
+        negative_prompt=full_negative_prompt,
+        num_inference_steps=steps,
+        guidance_scale=cfg_scale,
         width=width,
         height=height,
         generator=generator,
     ).images[0]
 
-    return image, seed
+    # 生成元数据
+    metadata = {
+        "prompt": full_prompt,
+        "negative_prompt": full_negative_prompt,
+        "base_model": BASE_MODEL,
+        "fixed_loras": [lora_id for lora_id, _ in FIXED_LORAS],
+        "optional_lora": optional_lora_name if optional_lora_name != "None" else None,
+        "lora_scale": lora_scale,
+        "seed": seed,
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "style": style,
+        "width": width,
+        "height": height,
+        "timestamp": datetime.datetime.now().isoformat()
+    }
 
+    # 生成文件名
+    timestamp = datetime.datetime.now().strftime("%y%m%d%H%M")
+    filename_base = f"{seed}-{timestamp}"
 
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
+    # 保存为WebP（高质量）
+    img_buffer = io.BytesIO()
+    image.save(img_buffer, format="WEBP", quality=95, method=6)
+    img_buffer.seek(0)
 
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
+    # 保存元数据为TXT
+    metadata_buffer = io.StringIO()
+    json.dump(metadata, metadata_buffer, indent=2, ensure_ascii=False)
+    metadata_buffer.seek(0)
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+    # 返回：图像、元数据、文件名
+    return (
+        image,
+        json.dumps(metadata, indent=2, ensure_ascii=False),
+        f"{filename_base}.webp",
+        f"{filename_base}.txt",
+        img_buffer.getvalue(),
+        metadata_buffer.getvalue().encode('utf-8')
+    )
 
-        result = gr.Image(label="Result", show_label=False)
+# ======================
+# Gradio UI
+# ======================
+with gr.Blocks(
+    theme=gr.themes.Soft(
+        primary_hue="blue",
+        secondary_hue="green",
+        neutral_hue="slate",
+    ).set(
+        body_background_fill="linear-gradient(135deg, #1e40af, #059669)",
+        button_primary_background_fill="white",
+        button_primary_text_color="#1e40af",
+        input_background_fill="rgba(255,255,255,0.9)",
+        text_size="lg",
+    ),
+    css="""
+    body { font-family: 'Helvetica Neue', 'Segoe UI', 'Arial', sans-serif; }
+    .gr-button { font-family: 'Helvetica Neue', 'Arial', sans-serif; font-weight: 500; }
+    .gr-textarea { font-family: 'Consolas', 'Monaco', 'Courier New', monospace; }
+    """,
+) as demo:
+    gr.Markdown(
+        """
+        # 🎨 AI Photo Generator (RealisticVision + LoRA)
+        **PRO + ZeroGPU Optimized | Multi-LoRA | Style Templates | Metadata Export**
+        """
+    )
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
+    with gr.Row():
+        with gr.Column(scale=3):
+            # a. 提示词输入框
+            prompt_input = gr.Textbox(
+                label="Prompt (Positive)",
+                placeholder="A beautiful woman, golden hour, soft sunlight...",
+                lines=5,
+                max_lines=20,
+                elem_classes=["gr-textarea"]
             )
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+            # b. 负提示词输入框
+            negative_prompt_input = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="blurry, low quality, deformed, cartoon, anime, text, watermark...",
+                lines=5,
+                max_lines=20,
+                elem_classes=["gr-textarea"]
             )
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            # c. 风格选择（单选）
+            style_radio = gr.Radio(
+                choices=list(STYLE_PROMPTS.keys()),
+                label="Style",
+                value="Realistic",
+                elem_classes=["gr-radio"]
+            )
 
+            # d. 种子选择
             with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                seed_input = gr.Slider(
+                    minimum=-1,
+                    maximum=99999999,
+                    step=1,
+                    value=DEFAULT_SEED,
+                    label="Seed (-1 = Random)"
                 )
+                seed_reset = gr.Button("Reset Seed")
 
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+            # e. 宽度选择
+            with gr.Row():
+                width_input = gr.Slider(
+                    minimum=512,
+                    maximum=1536,
+                    step=64,
+                    value=DEFAULT_WIDTH,
+                    label="Width"
+                )
+                width_reset = gr.Button("Reset Width")
+
+            # f. 高度选择
+            with gr.Row():
+                height_input = gr.Slider(
+                    minimum=512,
+                    maximum=1536,
+                    step=64,
+                    value=DEFAULT_HEIGHT,
+                    label="Height"
                 )
+                height_reset = gr.Button("Reset Height")
+
+            # g. LoRA选择（下拉）
+            optional_lora_dropdown = gr.Dropdown(
+                choices=list(OPTIONAL_LORA_MAP.keys()),
+                label="Optional LoRA",
+                value="None",
+                elem_classes=["gr-dropdown"]
+            )
 
+            # h. LoRA控制
             with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
+                lora_scale_slider = gr.Slider(
                     minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                    maximum=1.5,
+                    step=0.05,
+                    value=DEFAULT_LORA_SCALE,
+                    label="LoRA Scale"
                 )
+                lora_reset = gr.Button("Reset LoRA Scale")
 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
+            # i. 功能控制（Steps & CFG）
+            with gr.Row():
+                steps_slider = gr.Slider(
+                    minimum=10,
+                    maximum=100,
                     step=1,
-                    value=2,  # Replace with defaults that work for your model
+                    value=DEFAULT_STEPS,
+                    label="Steps"
+                )
+                cfg_slider = gr.Slider(
+                    minimum=1.0,
+                    maximum=20.0,
+                    step=0.5,
+                    value=DEFAULT_CFG,
+                    label="CFG Scale"
                 )
+                gen_reset = gr.Button("Reset Generation")
+
+            # m. 生成按钮
+            generate_btn = gr.Button("✨ Generate Image", variant="primary", size="lg")
+
+        with gr.Column(scale=2):
+            # j. 图片显示区
+            image_output = gr.Image(label="Generated Image", height=768, format="webp")
+
+            # k. 元数据显示区
+            metadata_output = gr.Textbox(
+                label="Metadata (JSON)",
+                lines=12,
+                max_lines=20,
+                elem_classes=["gr-textarea"]
+            )
 
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
+            # l. 下载按钮（并列）
+            with gr.Row():
+                download_img_btn = gr.Button("⬇️ Download Image (WebP)")
+                download_meta_btn = gr.Button("⬇️ Download Metadata (TXT)")
+
+            # 隐藏文件输出（用于下载）
+            hidden_img_file = gr.File(visible=False)
+            hidden_meta_file = gr.File(visible=False)
+
+    # ======================
+    # 事件绑定
+    # ======================
+
+    # 重置种子
+    seed_reset.click(fn=lambda: -1, outputs=seed_input)
+    # 重置宽度
+    width_reset.click(fn=lambda: DEFAULT_WIDTH, outputs=width_input)
+    # 重置高度
+    height_reset.click(fn=lambda: DEFAULT_HEIGHT, outputs=height_input)
+    # 重置LoRA缩放
+    lora_reset.click(fn=lambda: DEFAULT_LORA_SCALE, outputs=lora_scale_slider)
+    # 重置生成参数
+    gen_reset.click(
+        fn=lambda: (DEFAULT_STEPS, DEFAULT_CFG),
+        outputs=[steps_slider, cfg_slider]
+    )
+
+    # 生成
+    generate_btn.click(
+        fn=generate_image,
         inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
+            prompt_input, negative_prompt_input, style_radio,
+            seed_input, width_input, height_input,
+            optional_lora_dropdown, lora_scale_slider,
+            steps_slider, cfg_slider
         ],
-        outputs=[result, seed],
+        outputs=[
+            image_output, metadata_output,
+            hidden_img_file, hidden_meta_file,
+            hidden_img_file, hidden_meta_file
+        ]
+    )
+
+    # 下载图片
+    download_img_btn.click(
+        fn=None,
+        inputs=[hidden_img_file],
+        outputs=None,
+        js="(f) => { const a = document.createElement('a'); a.href = f; a.download = f.split('/').pop(); document.body.appendChild(a); a.click(); document.body.removeChild(a); }"
+    )
+
+    # 下载元数据
+    download_meta_btn.click(
+        fn=None,
+        inputs=[hidden_meta_file],
+        outputs=None,
+        js="(f) => { const a = document.createElement('a'); a.href = f; a.download = f.split('/').pop(); document.body.appendChild(a); a.click(); document.body.removeChild(a); }"
+    )
+
+    # 设置文件下载（通过返回值触发）
+    generate_btn.change(
+        fn=lambda img_bytes, meta_bytes, img_name, meta_name: (
+            gr.File(value=io.BytesIO(img_bytes), label=img_name, visible=True),
+            gr.File(value=io.BytesIO(meta_bytes), label=meta_name, visible=True)
+        ),
+        inputs=[hidden_img_file, hidden_meta_file, hidden_img_file, hidden_meta_file],
+        outputs=[hidden_img_file, hidden_meta_file]
     )
 
+# ======================
+# 启动
+# ======================
 if __name__ == "__main__":
     demo.launch()
+```