app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces
import os
import base64
import math

from PIL import Image
from diffusers import QwenImageEditPlusPipeline
from pillow_heif import register_heif_opener

from huggingface_hub import login
# from prompt_augment import PromptAugment
hf_token = os.getenv("hf")

space_id = os.getenv("SPACE_ID")
root_path = f"https://{space_id.replace('/', '-')}.hf.space" if space_id else ""

print(f"root path: {root_path}")

if hf_token:
    print("Secret loaded successfully.")
else:
    print("Secret not found. Check your Space settings.")
    
login(token=hf_token)

register_heif_opener()

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = QwenImageEditPlusPipeline.from_pretrained(
    "FireRedTeam/FireRed-Image-Edit-1.1", 
    torch_dtype=dtype
).to(device)
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()

#load lightning
pipe.load_lora_weights(
            "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo", 
            weight_name="FireRed-Image-Edit-Lightning-8steps-v1.0.safetensors", 
            adapter_name= "lightning"
)
pipe.load_lora_weights(
            "aiunivers/qwen-image-edit-plus-nsfw-lora", 
            weight_name="qwen-image-edit-plus-nsfw-lora.safetensors", 
            adapter_name= "aiuniverse"
)



# prompt_handler = PromptAugment()
LORA_REPO = "wiikoo/Qwen-lora-nsfw"
LORA_CONFIGS = {
    "CockQwen_v3":              "loras/CockQwen-v3.safetensors",
    "Eva_Qwen_V3":              "loras/Eva_Qwen_V3.safetensors",
    "Facial_Cumshots_V1":       "loras/Facial_Cumshots_For_Qwen_Image_V1.safetensors",
    "HearmemanAI_V3_Breasts":   "loras/HearmemanAI_V3_Rank64_BreastsLoRA_Epoch60.safetensors",
    "HearmemanAI_V4_Breasts":   "loras/HearmemanAI_V4_Rank128_BreastsLoRA_Epoch80.safetensors",
    "InniePussy":               "loras/InniePussy.safetensors",
    "JTT2_5":                   "loras/[QWEN] JTT2_5.safetensors",
    "LumiNude01a":              "loras/LumiNude01a_CE_QWEN_AIT3k.safetensors",
    "MEXX_QWEN_TG300":          "loras/MEXX_QWEN_TG300_23.safetensors",
    "Meta4":                    "loras/Meta4.safetensors",
    "MysticXXX":                "loras/Qwen-MysticXXX-v1.safetensors",
    "Nsfw_Body_V10":            "loras/Qwen_Nsfw_Body_V10-4K.safetensors",
    "Nsfw_Body_V14":            "loras/Qwen_Nsfw_Body_V14-10K.safetensors",
    "OilySkin_V2":              "loras/Oily Skin QWEN V2-GMR.safetensors",
    "PillowHump_2509":          "loras/PillowHump_2509.safetensors",
    "PutItHere_V2":             "loras/Put it here_Qwen edit_V2.0.safetensors",
    "PutItHere_V01":            "loras/put it here_QwenEdit_V0.1.safetensors",
    "Qwen4Play_v2":             "loras/Qwen4Play_v2.safetensors",
    "QwenHentai_v3":            "loras/QwenImageHentaiPIV_v3.1.safetensors",
    "Qwen_Helm":                "loras/Qwen-Image-Helm_v0.1.safetensors",
    "Qwen_NSFW_Beta1":          "loras/Qwen-NSFW.safetensors",
    "Qwen_NSFW_Beta2":          "loras/Qwen-NSFW-Beta2.safetensors",
    "Qwen_NSFW_Beta4":          "loras/Qwen-NSFW-Beta4.safetensors",
    "Qwen_NSFW_Beta5":          "loras/Qwen-NSFW-Beta5.safetensors",
    "Qwen_Real_Nud3s":          "loras/Qwen_Real_Nud3s.safetensors",
    "Qwen_Real_PS":             "loras/Qwen-Real PS_v1_83K.safetensors",
    "QwenSnofs_v1":             "loras/qwen_snofs.safetensors",
    "QwenSnofs_v1_1":           "loras/QwenSnofs1_1.safetensors",
    "Real_Breast_Nipples":      "loras/Real Breast Nipples-QWEN-[rbn]-GMR.safetensors",
    "SendDudes":                "loras/[QWEN] SendDudes.safetensors",
    "SendNudesLite":            "loras/SendNudesLite (Qwen).safetensors",
    "SendNudesPro_Beta":        "loras/[QWEN] Send Nudes Pro - Beta v1.safetensors",
    "Ultimate_Breast_Nipples":  "loras/Ultimate Realistic Breast NIPPLES-QWEN-[rab]-GMR.safetensors",
    "ass_up_QWEN":              "loras/ass_up_QWEN.safetensors",
    "barbell_nipples_QWEN":     "loras/QWEN_jtn_barbell.safetensors",
    "bfs_v2_face":              "loras-sfw/face_swap_5500_qwen_image_edit_2509_v1.safetensors",
    "bfs_v2_focus_face":        "loras-sfw/bfs_v2_000005000.safetensors",
    "big_nipples_QWEN":         "loras/big_nipples_QWEN.safetensors",
    "bumpynipples":             "loras/bumpynipples1.safetensors",
    "cmslt_cum_on_her":         "loras/cmslt_2509_2.safetensors",
    "consistence_edit_v1":      "loras-2/consistence_edit_v1.safetensors",
    "consistence_edit_v2":      "loras2/consistence_edit_v2.safetensors",
    "d33p7hroa7":               "loras/d33p7hroa7_qwen.safetensors",
    "d1ck_p3n1s_V1_1":          "loras/qwen-image_d!ck_P3N1S_LoRA_V1.1.safetensors",
    "goblin_anal_v1":           "loras/goblin_anal_v1_qwen.safetensors",
    "horseshoe_nipple_rings":   "loras/horseshoe_nipple_rings_QWEN.safetensors",
    "jib_nudity_fixer":         "loras/jib_qwen_fix_000002750.safetensors",
    "jillin":                   "loras/jillin1.safetensors",
    "male_nude":                "loras/lora_nudenan_v1.safetensors",
    "milk_juggs":               "loras/milk_juggs_QWEN.safetensors",
    "n00d_b":                   "loras/n00d-b-qwen.safetensors",
    "nsfw_adv_v1":              "loras/qwen-image_nsfw_adv_v1.0.safetensors",
    "p0ssy_lora_v1":            "loras/p0ssy_lora_v1.safetensors",
    "p3nis":                    "loras/p3nis.safetensors",
    "qwen_MCNL":                "loras/qwen_MCNL_v1.0.safetensors",
    "qwen_PENISLORA":           "loras/qwen-PENISLORA.safetensors",
    "qwen_hand_grab":           "loras/qwen_hand_grab_6000s.safetensors",
    "qwen_uncensor":            "loras/qwen_uncensor_000014928.safetensors",
    "reclining_nude":           "loras/reclining_nude_v1_000003500.safetensors",
    "remove_clothing":          "loras/qwen_image_edit_remove-clothing_v1.0.safetensors",
    "royal_treatment_V3":       "loras/royal+treatment+V3.safetensors",
    "sabi_character":           "loras-2/sabi_character_v1.safetensors",
    "snapchat_selfie":          "loras/qwen_image_snapchat.safetensors",
    "uka_qwen":                 "loras/uka_1_qwen.safetensors",
    "ultimate_realistic_breast":"loras/ultimate realistic breast.safetensors",
}

ADAPTER_SPECS = {
    "Covercraft": {
        "repo": "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo",
        "weights": "FireRed-Image-Edit-Covercraft.safetensors",
        "adapter_name": "covercraft",
    },
    "Lightning": {
        "repo": "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo",
        "weights": "FireRed-Image-Edit-Lightning-8steps-v1.0.safetensors",
        "adapter_name": "lightning",
    },
    "Makeup": {
        "repo": "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo",
        "weights": "FireRed-Image-Edit-Makeup.safetensors",
        "adapter_name": "makeup",
    }
}

LOADED_ADAPTERS = set()
LORA_OPTIONS = ["None"] + list(LORA_CONFIGS.keys())
print(f"lora option: {LORA_OPTIONS[0]}")


def load_lora(lora_name):
    """加载并激活指定的 LoRA"""
    if lora_name == "None" or not lora_name:
        if LOADED_ADAPTERS:
            pipe.set_adapters([], adapter_weights=[])
        return

    
    # spec = ADAPTER_SPECS.get(lora_name)
    # if not spec:
    #     raise gr.Error(f"LoRA 配置未找到: {lora_name}")

    # adapter_name = spec["adapter_name"]

    if lora_name not in LOADED_ADAPTERS:
        print(f"--- Downloading and Loading Adapter: {lora_name} ---")

        if lora_name == "Lightning":
            pipe.load_lora_weights(
                "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo", 
                weight_name="FireRed-Image-Edit-Lightning-8steps-v1.0.safetensors", 
                adapter_name= lora_name
            )
            LOADED_ADAPTERS.add(lora_name)

        else:
            try:
                pipe.load_lora_weights(
                    LORA_REPO, 
                    weight_name=LORA_CONFIGS[lora_name], 
                    adapter_name= lora_name
                )
                LOADED_ADAPTERS.add(lora_name)
            except Exception as e:
                raise gr.Error(f"Failed to load adapter {lora_name}: {e}")
    else:
        print(f"--- Adapter {lora_name} is already loaded ---")

    pipe.set_adapters(["lightning", "aiuniverse", lora_name], adapter_weights=[1.0, 1.0, 0.5])


MAX_SEED = np.iinfo(np.int32).max
MAX_INPUT_IMAGES = 3


def limit_images(images):
    if images is None:
        return None
    if len(images) > MAX_INPUT_IMAGES:
        gr.Info(f"最多支持 {MAX_INPUT_IMAGES} 张图片，已自动移除多余图片")
        return images[:MAX_INPUT_IMAGES]
    return images


def calculate_dimensions(target_area, ratio):
    width = math.sqrt(target_area * ratio)
    height = width / ratio
    width = round(width / 32) * 32
    height = round(height / 32) * 32
    return int(width), int(height)


def update_dimensions_on_upload(images, max_area=1024*1024):
    if images is None or len(images) == 0:
        return 0, 0
    
    try:
        first_item = images[0]
        if isinstance(first_item, tuple):
            img = first_item[0]
        else:
            img = first_item
        
        if isinstance(img, Image.Image):
            pil_img = img
        elif isinstance(img, str):
            pil_img = Image.open(img)
        else:
            return 0, 0
        
        h, w = pil_img.height, pil_img.width
        is_multi_image = len(images) > 1
        
        if not is_multi_image:
            return 0, 0
        
        ratio = w / h
        new_w, new_h = calculate_dimensions(max_area, ratio)
        return new_h, new_w
    except Exception as e:
        print(f"获取图片尺寸失败: {e}")
        return 0, 0


@spaces.GPU(duration=180)
def infer(
    input_images,
    prompt,
    lora_choice,
    seed=42,
    true_guidance_scale=4.0,
    num_inference_steps=8,
    height=None,
    width=None,
    # rewrite_prompt=False,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):
    negative_prompt = " "

    seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    

    load_lora(lora_choice)
    pil_images = []
    if input_images is not None:
        for item in input_images[:MAX_INPUT_IMAGES]:
            try:
                if isinstance(item, tuple):
                    img = item[0]
                else:
                    img = item
                    
                if isinstance(img, Image.Image):
                    pil_images.append(img.convert("RGB"))
                elif isinstance(img, str):
                    pil_images.append(Image.open(img).convert("RGB"))
            except Exception as e:
                print(f"处理图片出错: {e}")
                continue
    
    if height == 0:
        height = None
    if width == 0:
        width = None
    
    # if rewrite_prompt and len(pil_images) > 0:
    #     # prompt = prompt_handler.predict(prompt, [pil_images[0]])
    #     print(f"Rewritten Prompt: {prompt}")
    
    if pil_images:
        for i, img in enumerate(pil_images):
            print(f"    [{i}] size: {img.width}x{img.height}")

    num_inference_steps = max(1, int(num_inference_steps or 1))
    
    images = pipe(
        image=pil_images if len(pil_images) > 0 else None,
        prompt=prompt,
        height=height,
        width=width,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=1.0,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=num_images_per_prompt,
    ).images

    return images, seed

css = """
#col-container { margin: 0 auto; max-width: 1200px; }
#edit-btn { height: 100% !important; min-height: 42px; }
"""



def get_image_base64(image_path):
    with open(image_path, "rb") as img_file:
        return base64.b64encode(img_file.read()).decode('utf-8')


logo_base64 = None

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):

        gr.Markdown(f"Supports multi-image input (up to {MAX_INPUT_IMAGES} images.)")
        
        with gr.Row():
            with gr.Column(scale=1):
                input_images = gr.Gallery(
                    label="Upload Images",
                    type="pil",
                    interactive=True,
                    height=300,
                    columns=3,
                    object_fit="contain",
                )
            
            with gr.Column(scale=1):
                result = gr.Gallery(
                    label="Output Images",
                    type="pil",
                    height=300,
                    columns=2,
                    object_fit="contain",
                )
        
        prompt = gr.Textbox(
            label="Edit Prompt",
            placeholder="e.g., transform into anime..",
        )
        
        with gr.Row(equal_height=True):
            with gr.Column(scale=5):
                lora_choice = gr.Dropdown(
                    label="Choose Lora",
                    choices=LORA_OPTIONS,
                    value=LORA_OPTIONS[0] if LORA_OPTIONS else "None",
                )
            with gr.Column(scale=4):
                run_button = gr.Button("Edit Image", variant="primary", elem_id="edit-btn")

        with gr.Accordion("Advanced Settings", open=True):
            with gr.Row():
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
            
            
            with gr.Row():
                true_guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=4.0)
                num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=30, step=1, value=8)
            
            with gr.Row():
                height = gr.Slider(label="Height (0=auto)", minimum=0, maximum=2048, step=8, value=0)
                width = gr.Slider(label="Width (0=auto)", minimum=0, maximum=2048, step=8, value=0)
            
            with gr.Row():
                # rewrite_prompt = gr.Checkbox(label="Rewrite Prompt", value=True)
                num_images_per_prompt = gr.Slider(label="Num Images", minimum=1, maximum=4, step=1, value=1)


    # 监听 LoRA 选择变化：Lightning 时锁定参数
    def on_lora_change(lora_name):
        
        return (
            gr.update(value=8, interactive=False),      # num_inference_steps
            gr.update(value=1.0, interactive=False),    # true_guidance_scale
            gr.update(value=43, interactive=True),      # seed
            # gr.update(value=False, interactive=False),  # randomize_seed
        )


    lora_choice.change(
        fn=on_lora_change,
        inputs=[lora_choice],
        outputs=[num_inference_steps, true_guidance_scale, seed],
    )
    def on_image_upload(images):
        limited = limit_images(images)
        h, w = update_dimensions_on_upload(limited)
        return limited, h, w
    
    input_images.upload(
        fn=on_image_upload,
        inputs=[input_images],
        outputs=[input_images, height, width],
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            input_images,
            prompt, lora_choice, seed,
            true_guidance_scale, num_inference_steps,
            height, width, num_images_per_prompt,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch(css=css, allowed_paths=["./"])
    # demo.launch(allowed_paths=["./"], ssr_mode=False, css=css, root_path=root_path, 
    #             show_error=True, auth=None)