app.py

# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gradio as gr
import spaces
import torch
import os
from PIL import Image
from diffusers.utils import load_image

# 导入你的两个 Pipeline
from dreamlite import DreamLitePipeline 
from dreamlite import DreamLiteMobilePipeline 

from huggingface_hub import snapshot_download

from transformers import CLIPFeatureExtractor
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
import numpy as np


# ==========================================
# 1. 全局配置与模型缓存管理
# ==========================================
HF_TOKEN = os.environ.get("HF_TOKEN")
device = "cpu"
dtype = torch.bfloat16

# 加载 safety checker（启动时加载一次）
safety_checker = StableDiffusionSafetyChecker.from_pretrained(
    "CompVis/stable-diffusion-safety-checker"
)
feature_extractor = CLIPFeatureExtractor.from_pretrained(
    "openai/clip-vit-base-patch32"
)

def check_safety(image):
    """检查生成的图像是否安全，返回 (image, is_nsfw)"""
    safety_input = feature_extractor(images=image, return_tensors="pt")
    np_image = np.array(image)
    _, has_nsfw = safety_checker(
        images=[np_image],
        clip_input=safety_input.pixel_values
    )
    return has_nsfw[0]

base_path = snapshot_download("carlofkl/DreamLite-base", token=HF_TOKEN)
mobile_path = snapshot_download("carlofkl/DreamLite-mobile", token=HF_TOKEN)

# 预加载两个模型到显存中，避免推理时动态加载
print("Loading DreamLite-base...")
pipe_base = DreamLitePipeline.from_pretrained(base_path, torch_dtype=dtype).to(device)
print("DreamLite-base Loaded Successfully!")

print("Loading DreamLite-mobile...")
pipe_mobile = DreamLiteMobilePipeline.from_pretrained(mobile_path, torch_dtype=dtype).to(device)
print("DreamLite-mobile Loaded Successfully!")

MODEL_CONFIGS = {
    "DreamLite-base": pipe_base,
    "DreamLite-mobile": pipe_mobile,
}

BASE_RESOLUTIONS = [
    "1024 × 1024 (1:1)",
    "1152 × 896 (9:7)",
    "896 × 1152 (7:9)",
    "1216 × 832 (3:2)",
    "832 × 1216 (2:3)",
    "1344 × 768 (16:9)",
    "768 × 1344 (9:16)",
]

def parse_resolution(res_str):
    """从分辨率字符串中解析宽高，例如 '1024 × 1024 (1:1)' -> (1024, 1024)"""
    parts = res_str.split("(")[0].strip().split("×")
    w = int(parts[0].strip())
    h = int(parts[1].strip())
    return w, h

# ==========================================
# 2. 定义推理函数
# ==========================================
@spaces.GPU
def generate_image(
    model_choice,
    prompt,
    image,
    resolution,
    num_inference_steps,
    guidance_scale,
    image_guidance_scale,
    seed
):
    # 从预加载的模型中选择
    pipe = MODEL_CONFIGS[model_choice]

    # 强制将种子转为 Tensor Generator 以保证可复现
    generator = torch.Generator(device="cpu").manual_seed(seed)
    
    # 将 Gradio 传入的图片 (如果有的话) 转换为 PIL 格式
    input_image = image if image is not None else None

    if model_choice == "DreamLite-base":
        width, height = parse_resolution(resolution)
    else:
        # Mobile 版本固定 1024x1024
        width, height = parse_resolution(resolution)
    
    if image is not None: width, height = image.size
    
    # 调用对应的 Pipeline
    out = pipe(
        prompt=prompt,
        image=input_image,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        image_guidance_scale=image_guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
    ).images[0]

    if out.size != (width, height):
        out = out.resize((width, height), resample=Image.LANCZOS)

    # Safety Check
    if check_safety(out):
        raise gr.Error("⚠️ The generated image has been blocked by our safety filter. "
                       "Please try a different prompt.")
    
    return out

# ==========================================
# 3. UI 联动：切换模型时更新参数面板
# ==========================================
def on_model_change(model_choice):
    """
    切换模型时自动调整 UI 组件：
    - Base: 显示分辨率选择，默认 28 步
    - Mobile: 隐藏分辨率选择（固定 1024×1024），默认 4 步
    """
    if model_choice == "DreamLite-base":
        return (
            gr.update(visible=True, value="1024 × 1024 (1:1)"),  # 分辨率选择可见
            gr.update(value=28),                                   # 默认 28 步
            gr.update(value=3.5),                                  # guidance scale
        )
    else:
        return (
            gr.update(visible=True, value="1024 × 1024 (1:1)"),    # 分辨率选择隐藏
            gr.update(value=4),                                    # 默认 4 步
            gr.update(value=1.0),                                  # guidance scale
        )

# ==========================================
# 4. 搭建 Gradio 页面
# ==========================================
with gr.Blocks(title="DreamLite Demo") as demo:
    gr.Markdown("# 🌟 DreamLite: Efficient On-Device Generation and Editing")
    gr.Markdown("Select a model version, then generate images from text or upload an image to edit it based on instructions.")
    
    with gr.Row():
        with gr.Column():
            # 新增：模型选择下拉框
            model_dropdown = gr.Dropdown(
                choices=list(MODEL_CONFIGS.keys()),
                value="DreamLite-base", # 默认选中 base
                label="Select Model Version",
                interactive=True
            )
            
            # 输入组件
            prompt_input = gr.Textbox(label="Prompt / Instruction", placeholder="e.g. A photo of a dog...", lines=3)
            image_input = gr.Image(type="pil", label="Input Image (Optional for Editing)")

            # 分辨率选择（仅 Base 版本可见）
            resolution_dropdown = gr.Dropdown(
                choices=BASE_RESOLUTIONS,
                value="1024 × 1024 (1:1)",
                label="Resolution (Base model only, Mobile fixed at 1024×1024)",
                interactive=True,
                visible=True
            )
            
            with gr.Accordion("Advanced Options", open=False):
                steps_slider = gr.Slider(minimum=1, maximum=50, value=28, step=1, label="Inference Steps")
                guidance_slider = gr.Slider(minimum=0.0, maximum=20.0, value=3.5, step=0.1, label="Guidance Scale")
                img_guidance_slider = gr.Slider(minimum=0.0, maximum=5.0, value=1.0, step=0.1, label="Image Guidance Scale")
                seed_slider = gr.Slider(minimum=0, maximum=999999, value=42, step=1, label="Seed")
                
            submit_btn = gr.Button("Generate / Edit", variant="primary")
            
        with gr.Column():
            # 输出组件
            output_image = gr.Image(type="pil", label="Output Image")
            
    # 模型切换时联动更新 UI
    model_dropdown.change(
        fn=on_model_change,
        inputs=[model_dropdown],
        outputs=[resolution_dropdown, steps_slider, guidance_slider]
    )

    # 绑定点击事件 (注意 inputs 列表增加了 model_dropdown 作为第一个参数)
    submit_btn.click(
        fn=generate_image,
        inputs=[model_dropdown, prompt_input, image_input, resolution_dropdown, steps_slider, guidance_slider, img_guidance_slider, seed_slider],
        outputs=[output_image]
    )
    
    # 示例区 (同步加上对应的模型选择)
    gr.Examples(
        examples=[
            ["DreamLite-base", "A close-up of a fire spitting dragon, cinematic shot.", None, "1216 × 832 (3:2)", 28, 3.5, 1.0, 123],
            ["DreamLite-mobile", "A portrait of a young woman with flowers.", None, "1024 × 1024 (1:1)", 4, 3.5, 1.0, 42],
            ["DreamLite-mobile", "Make it look like a pencil sketch", "assets/example.png", "1024 × 1024 (1:1)", 4, 1.0, 1.0, 42],
        ],
        inputs=[model_dropdown, prompt_input, image_input, resolution_dropdown, steps_slider, guidance_slider, img_guidance_slider, seed_slider]
    )

# ==========================================
# 5. 启动应用
# ==========================================
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)