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)