Upload nag_multi_app.py

nag_multi_app.py

@@ -0,0 +1,401 @@
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import random
+from PIL import Image
+import math
+
+# --- nag_app.pyから移植した機能 ---
+# 翻訳ライブラリのインポート
+from deep_translator import GoogleTranslator
+from langdetect import detect
+
+# NAG対応パイプラインのインポート
+# 注: このコードを実行するには、nag_app.pyのHugging Face Spaceから
+# `src`ディレクトリ（pipeline_flux_kontext_nag.pyとtransformer_flux.pyを含む）を
+# このファイルと同じ階層に配置する必要があります。
+from src.pipeline_flux_kontext_nag import NAGFluxKontextPipeline
+from src.transformer_flux import NAGFluxTransformer2DModel
+# --- ここまでが移植部分 ---
+
+# エラー解決のためにdiffusersの内部マッピングをインポート
+from diffusers.loaders.peft import _SET_ADAPTER_SCALE_FN_MAPPING
+
+
+# 定数の設定
+MAX_SEED = np.iinfo(np.int32).max
+DEFAULT_NAG_NEGATIVE_PROMPT = "Low resolution, blurry, lack of details, big head"
+OUTPUT_RESOLUTION = 1024
+
+# --- nag_app.pyから移植したモデル読み込み処理 ---
+# NAG対応のKontextモデルをロード
+transformer = NAGFluxTransformer2DModel.from_pretrained(
+    "black-forest-labs/FLUX.1-Kontext-dev",
+    subfolder="transformer",
+    torch_dtype=torch.bfloat16,
+)
+pipe = NAGFluxKontextPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Kontext-dev",
+    transformer=transformer,
+    torch_dtype=torch.bfloat16,
+)
+pipe = pipe.to("cuda")
+# --- ここまでが移植部分 ---
+
+# --- LoRAの読み込み処理 (5つ) ---
+print("Loading LoRA weights...")
+# LoRA名とアダプター名のマッピング
+LORA_MAPPING = {
+    "Hyper-SD": "hyper",
+    "Relighting": "relight",
+    "LoRA 3": "lora_3",
+    "LoRA 4": "lora_4",
+    "LoRA 5": "lora_5",
+}
+# 1. Hyper-SD LoRA
+pipe.load_lora_weights(
+    "ByteDance/Hyper-SD",
+    weight_name="Hyper-FLUX.1-dev-8steps-lora.safetensors",
+    adapter_name=LORA_MAPPING["Hyper-SD"]
+)
+# 2. Relighting LoRA
+pipe.load_lora_weights(
+    "linoyts/relighting-kontext-dev-lora",
+    weight_name="relighting-kontext-dev-lora.safetensors",
+    adapter_name=LORA_MAPPING["Relighting"]
+)
+# 3. 追加のLoRA 3 (後で設定)
+# ★ 注意: 以下のリポジトリ名とファイル名は仮のものです。後で正しいものに置き換えてください。
+try:
+    pipe.load_lora_weights(
+        "author/repo_name_3", # 例: "cagliostrolab/animagine-xl-3.0"
+        weight_name="lora_file_3.safetensors", # 例: "animagine-xl-3.0.safetensors"
+        adapter_name=LORA_MAPPING["LoRA 3"]
+    )
+except Exception as e:
+    print(f"Warning: Could not load {list(LORA_MAPPING.keys())[2]}. Please check repository and file names. Error:", e)
+
+# 4. 追加のLoRA 4 (後で設定)
+try:
+    pipe.load_lora_weights(
+        "author/repo_name_4",
+        weight_name="lora_file_4.safetensors",
+        adapter_name=LORA_MAPPING["LoRA 4"]
+    )
+except Exception as e:
+    print(f"Warning: Could not load {list(LORA_MAPPING.keys())[3]}. Please check repository and file names. Error:", e)
+
+# 5. 追加のLoRA 5 (後で設定)
+try:
+    pipe.load_lora_weights(
+        "author/repo_name_5",
+        weight_name="lora_file_5.safetensors",
+        adapter_name=LORA_MAPPING["LoRA 5"]
+    )
+except Exception as e:
+    print(f"Warning: Could not load {list(LORA_MAPPING.keys())[4]}. Please check repository and file names. Error:", e)
+
+print("LoRA weights loading process finished.")
+# --- ここまでが変更部分 ---
+
+# カスタムモデルをdiffusersのLoRA対応表に登録する
+_SET_ADAPTER_SCALE_FN_MAPPING[NAGFluxTransformer2DModel.__name__] = _SET_ADAPTER_SCALE_FN_MAPPING["FluxTransformer2DModel"]
+print("Custom model 'NAGFluxTransformer2DModel' registered for LoRA.")
+
+
+def round_to_multiple(number, multiple=8):
+    return multiple * round(number / multiple)
+
+def concatenate_images(images, direction="horizontal"):
+    if not images: return None
+    valid_images = [img for img in images if img is not None]
+    if not valid_images: return None
+    if len(valid_images) == 1: return valid_images[0].convert("RGB")
+    valid_images = [img.convert("RGB") for img in valid_images]
+    if direction == "horizontal":
+        total_width = sum(img.width for img in valid_images)
+        max_height = max(img.height for img in valid_images)
+        concatenated = Image.new('RGB', (total_width, max_height), (255, 255, 255))
+        x_offset = 0
+        for img in valid_images:
+            y_offset = (max_height - img.height) // 2
+            concatenated.paste(img, (x_offset, y_offset))
+            x_offset += img.width
+    else:
+        max_width = max(img.width for img in valid_images)
+        total_height = sum(img.height for img in valid_images)
+        concatenated = Image.new('RGB', (max_width, total_height), (255, 255, 255))
+        y_offset = 0
+        for img in valid_images:
+            x_offset = (max_width - img.width) // 2
+            concatenated.paste(img, (x_offset, y_offset))
+            y_offset += img.height
+    return concatenated
+
+@spaces.GPU(duration=25)
+# ★ infer関数の引数に negative_prompt を追加
+def infer(input_images, prompt, negative_prompt, seed, randomize_seed, guidance_scale, nag_negative_prompt, nag_scale, num_inference_steps,
+          # LoRAの有効/無効と強度を個別に受け取る
+          enable_lora1, weight_lora1,
+          enable_lora2, weight_lora2,
+          enable_lora3, weight_lora3,
+          enable_lora4, weight_lora4,
+          enable_lora5, weight_lora5,
+          progress=gr.Progress(track_tqdm=True)):
+    
+    active_adapters = []
+    active_weights = []
+    
+    lora_params = [
+        (enable_lora1, weight_lora1, "Hyper-SD"),
+        (enable_lora2, weight_lora2, "Relighting"),
+        (enable_lora3, weight_lora3, "LoRA 3"),
+        (enable_lora4, weight_lora4, "LoRA 4"),
+        (enable_lora5, weight_lora5, "LoRA 5"),
+    ]
+    
+    for is_enabled, weight, name in lora_params:
+        if is_enabled:
+            adapter_name = LORA_MAPPING[name]
+            active_adapters.append(adapter_name)
+            active_weights.append(weight)
+            print(f"Applying {name} LoRA with weight {weight}")
+
+    if active_adapters:
+        pipe.set_adapters(active_adapters, adapter_weights=active_weights)
+    else:
+        print("No LoRA selected. Running without LoRA.")
+        pipe.disable_lora()
+
+    prompt = prompt.strip()
+    if prompt:
+        print(f"Original prompt: {prompt}")
+        try:
+            detected_lang = detect(prompt)
+            if detected_lang != 'en':
+                print(f"Detected language: {detected_lang}. Translating to English...")
+                translated_prompt = GoogleTranslator(source=detected_lang, target='en').translate(prompt)
+                prompt = translated_prompt
+                print(f"Translated prompt: {prompt}")
+            else:
+                print("Prompt is already in English.")
+        except Exception as e:
+            print(f"Warning: Translation or language detection failed: {e}. Using original prompt.")
+
+    # ★ negative_promptを処理するコードを追加
+    negative_prompt = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    if input_images is None:
+        raise gr.Error("Please upload at least one image.")
+    
+    if not isinstance(input_images, list):
+        input_images = [input_images]
+    
+    valid_images = [img[0] for img in input_images if img is not None]
+    
+    if not valid_images:
+        raise gr.Error("Please upload at least one valid image.")
+
+    if len(valid_images) == 1:
+        print("Single image detected. Calculating aspect-ratio aware dimensions.")
+        input_for_pipe = valid_images[0]
+        
+        input_width, input_height = input_for_pipe.size
+        aspect_ratio = input_width / input_height
+        target_pixels = OUTPUT_RESOLUTION * OUTPUT_RESOLUTION
+        
+        final_height = int(math.sqrt(target_pixels / aspect_ratio))
+        final_width = int(aspect_ratio * final_height)
+        
+        final_width = round_to_multiple(final_width, 8)
+        final_height = round_to_multiple(final_height, 8)
+        
+        print(f"Output dimensions set to: {final_width}x{final_height}")
+
+    else:
+        print(f"Multiple ({len(valid_images)}) images detected. Using fixed 1024x1024 output.")
+        input_for_pipe = concatenate_images(valid_images, "horizontal")
+        if input_for_pipe is None:
+            raise gr.Error("Failed to process the input images.")
+        
+        final_width = OUTPUT_RESOLUTION
+        final_height = OUTPUT_RESOLUTION
+    
+    final_prompt = f"From the provided reference images, create a unified, cohesive image such that {prompt}. Maintain the identity and characteristics of each subject while adjusting their proportions, scale, and positioning to create a harmonious, naturally balanced composition. Blend and integrate all elements seamlessly with consistent lighting, perspective, and style.the final result should look like a single naturally captured scene where all subjects are properly sized and positioned relative to each other, not assembled from multiple sources."
+    
+    # ★ pipe()呼び出しに negative_prompt を追加
+    image = pipe(
+        image=input_for_pipe,
+        prompt=final_prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        nag_negative_prompt=nag_negative_prompt,
+        nag_scale=nag_scale,
+        width=final_width,
+        height=final_height,
+        num_inference_steps=num_inference_steps,
+        generator=torch.Generator().manual_seed(seed),
+    ).images[0]
+    
+    pipe.disable_lora()
+    
+    return image, seed, gr.update(visible=True)
+
+css="""
+#col-container {
+    margin: 0 auto;
+    max-width: 960px;
+}
+.lora-row {
+    align-items: center;
+    margin-bottom: 8px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""# FLUX.1 Kontext [dev] - Multi-Image with NAG
+        Compose a new image from multiple images using FLUX.1 Kontext, enhanced with Normalized Attention Guidance (NAG) and automatic prompt translation.
+        - **Single Image Input**: Output will match the input aspect ratio.
+        - **Multiple Image Inputs**: Output will be a fixed 1024x1024 resolution.
+        """)
+        with gr.Row():
+            with gr.Column():
+                input_images = gr.Gallery(
+                    label="Upload image(s) for editing", 
+                    show_label=True,
+                    elem_id="gallery_input",
+                    columns=3,
+                    rows=2,
+                    object_fit="contain",
+                    height="auto",
+                    file_types=['image'],
+                    type='pil'
+                )
+                
+                with gr.Row():
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="Enter your prompt (auto-translates to English)",
+                        container=False,
+                    )
+                    run_button = gr.Button("Run", scale=0)
+                    
+                with gr.Accordion("Advanced Settings", open=False):
+                    # --- ★ UIを修正: 各LoRAコンポーネントを個別の変数として定義 ---
+                    gr.Markdown("### LoRA Settings")
+                    
+                    with gr.Row(elem_classes="lora-row"):
+                        enable_lora1 = gr.Checkbox(label="Hyper-SD", value=True, scale=1)
+                        weight_lora1 = gr.Slider(label="Weight", minimum=0.0, maximum=2.0, step=0.02, value=0.12, scale=3, visible=True)
+                    
+                    with gr.Row(elem_classes="lora-row"):
+                        enable_lora2 = gr.Checkbox(label="Relighting", value=False, scale=1)
+                        weight_lora2 = gr.Slider(label="Weight", minimum=0.0, maximum=2.0, step=0.05, value=1.0, scale=3, visible=False)
+
+                    with gr.Row(elem_classes="lora-row"):
+                        enable_lora3 = gr.Checkbox(label="LoRA 3", value=False, scale=1)
+                        weight_lora3 = gr.Slider(label="Weight", minimum=0.0, maximum=2.0, step=0.05, value=0.8, scale=3, visible=False)
+
+                    with gr.Row(elem_classes="lora-row"):
+                        enable_lora4 = gr.Checkbox(label="LoRA 4", value=False, scale=1)
+                        weight_lora4 = gr.Slider(label="Weight", minimum=0.0, maximum=2.0, step=0.05, value=0.8, scale=3, visible=False)
+
+                    with gr.Row(elem_classes="lora-row"):
+                        enable_lora5 = gr.Checkbox(label="LoRA 5", value=False, scale=1)
+                        weight_lora5 = gr.Slider(label="Weight", minimum=0.0, maximum=2.0, step=0.05, value=0.8, scale=3, visible=False)
+                    # --- ★ ここまでが変更部分 ---
+
+                    gr.Markdown("### Generation Settings")
+                    
+                    # ★ UIに negative_prompt を追加
+                    negative_prompt = gr.Text(
+                        label="Negative Prompt (Standard)",
+                        placeholder="Enter concepts to avoid (e.g., ugly, deformed)",
+                        max_lines=2,
+                    )
+                    
+                    num_inference_steps = gr.Slider(
+                        label="Inference Steps",
+                        minimum=8,
+                        maximum=50,
+                        step=1,
+                        value=8,
+                    )
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=1,
+                        maximum=10,
+                        step=0.25,
+                        value=4.5,
+                    )
+                    nag_negative_prompt = gr.Text(
+                        label="Negative Prompt for NAG",
+                        value=DEFAULT_NAG_NEGATIVE_PROMPT,
+                        max_lines=2,
+                        placeholder="Enter concepts to avoid with NAG",
+                    )
+                    nag_scale = gr.Slider(
+                        label="NAG Scale",
+                        minimum=0.0,
+                        maximum=20.0,
+                        step=0.25,
+                        value=3.5
+                    )
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0,
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                    
+            with gr.Column():
+                result = gr.Image(label="Result", show_label=False, interactive=False, format="png")
+                reuse_button = gr.Button("Reuse this image", visible=False)
+    
+    # ★ イベントハンドラを更新: all_inputsに negative_prompt を追加
+    all_inputs = [
+        input_images, prompt, negative_prompt, seed, randomize_seed, guidance_scale, 
+        nag_negative_prompt, nag_scale, num_inference_steps,
+        enable_lora1, weight_lora1,
+        enable_lora2, weight_lora2,
+        enable_lora3, weight_lora3,
+        enable_lora4, weight_lora4,
+        enable_lora5, weight_lora5,
+    ]
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn = infer,
+        inputs = all_inputs,
+        outputs = [result, seed, reuse_button]
+    )
+    # --- ★ ここまでが変更部分 ---
+    
+    reuse_button.click(
+        fn = lambda image: [image] if image is not None else [],
+        inputs = [result],
+        outputs = [input_images]
+    )
+
+    # --- ★ 各チェックボックスとスライダーの表示を個別に連動させる ---
+    def update_visibility(is_checked):
+        return gr.update(visible=is_checked)
+
+    enable_lora1.change(fn=update_visibility, inputs=enable_lora1, outputs=weight_lora1)
+    enable_lora2.change(fn=update_visibility, inputs=enable_lora2, outputs=weight_lora2)
+    enable_lora3.change(fn=update_visibility, inputs=enable_lora3, outputs=weight_lora3)
+    enable_lora4.change(fn=update_visibility, inputs=enable_lora4, outputs=weight_lora4)
+    enable_lora5.change(fn=update_visibility, inputs=enable_lora5, outputs=weight_lora5)
+    # --- ★ ここまでが変更部分 ---
+
+demo.launch()