app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces
from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler
from optimization import optimize_pipeline_
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
import math
import os

import os
import spaces
import torch
from diffusers import AutoencoderKLWan, WanPipeline, WanImageToVideoPipeline, UniPCMultistepScheduler
from diffusers.utils import export_to_video
import gradio as gr
import tempfile
from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image
import random


HF_MODEL = os.environ.get("HF_UPLOAD_REPO", "rahul7star/qwen-edit-img-repo")

# --- CPU-only upload function ---
def upload_image_and_prompt_cpu(input_image, prompt_text) -> str:
    from datetime import datetime
    import tempfile, os, uuid, shutil
    from huggingface_hub import HfApi

    # Instantiate the HfApi class
    api = HfApi()
    print(prompt_text)

    today_str = datetime.now().strftime("%Y-%m-%d")
    unique_subfolder = f"Upload-Image-{uuid.uuid4().hex[:8]}"
    hf_folder = f"{today_str}/{unique_subfolder}"

    # Save image temporarily
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_img:
        if isinstance(input_image, str):
            shutil.copy(input_image, tmp_img.name)
        else:
            input_image.save(tmp_img.name, format="PNG")
        tmp_img_path = tmp_img.name

    # Upload image using HfApi instance
    api.upload_file(
        path_or_fileobj=tmp_img_path,
        path_in_repo=f"{hf_folder}/input_image.png",
        repo_id=HF_MODEL,
        repo_type="model",
        token=os.environ.get("HUGGINGFACE_HUB_TOKEN")
    )

    # Save prompt as summary.txt
    summary_file = tempfile.NamedTemporaryFile(delete=False, suffix=".txt").name
    with open(summary_file, "w", encoding="utf-8") as f:
        f.write(prompt_text)

    api.upload_file(
        path_or_fileobj=summary_file,
        path_in_repo=f"{hf_folder}/summary.txt",
        repo_id=HF_MODEL,
        repo_type="model",
        token=os.environ.get("HUGGINGFACE_HUB_TOKEN")
    )

    # Cleanup
    os.remove(tmp_img_path)
    os.remove(summary_file)

    return hf_folder

# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Scheduler configuration for Lightning
scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}

# Initialize scheduler
scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)

# Load model
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509",
    scheduler=scheduler,
    torch_dtype=dtype
).to(device)

pipe.load_lora_weights(
    "rahul7star/qwen-char-lora",
    weight_name="qwen_lora/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16_dim1.safetensors"
)
pipe.fuse_lora(lora_scale=1.0)

# pipe.load_lora_weights(
#     "rahul7star/qwen-char-lora",
#     weight_name="qwen_lora/qwen-multiple-angle.safetensors",
# )
# pipe.fuse_lora(lora_scale=1.0)

pipe.load_lora_weights(
    "rahul7star/qwen-char-lora",
    weight_name="qwen_lora/qwen-multiple-char.safetensors",
)
pipe.fuse_lora(lora_scale=1.0)


pipe.transformer.__class__ = QwenImageTransformer2DModel
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")

# --- Constants ---
MAX_SEED = np.iinfo(np.int32).max
PROMPTS = {
    "front": "Move the camera to a front-facing position so the full body of the character is visible. The character stands with both arms extended slightly downward and close to the thighs, keeping the body evenly balanced on both sides. The legs are positioned symmetrically with a narrow stance. The background is plain white.",
    "back": "Move the camera to a back-facing position so the full body of the character is visible. Background is plain white.",
    "left": "Move the camera to a side view (profile) from the left so the full body of the character is visible. Background is plain white.",
    "right": "Move the camera to a side view (profile) from the right so the full body of the character is visible. Background is plain white."
}

# NEW: 出力解像度プリセット
RESOLUTIONS = {
    "1:4": (512, 2048),
    "1:3": (576, 1728),
    "nealy 9:16": (768, 1344),
    "nealy 2:3": (832, 1216),
    "3:4": (896, 1152),
}

def _append_prompt(base: str, extra: str) -> str:
    extra = (extra or "").strip()
    return (base if not extra else f"{base} {extra}").strip()

def generate_single_view(input_images, prompt, seed, num_inference_steps, true_guidance_scale):
    generator = torch.Generator(device=device).manual_seed(seed)
    print(prompt)
    
    result = pipe(
        image=input_images if input_images else None,
        prompt=prompt,
        negative_prompt=" ",
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=1,
    ).images
    try:
        upload_image_and_prompt_cpu(result[0], prompt)
    except Exception as e:
        print("Upload failed:", e)
    return result[0]

def concat_images_horizontally(images, bg_color=(255, 255, 255)):
    images = [img.convert("RGB") for img in images if img is not None]
    if not images:
        return None
    h = max(img.height for img in images)
    resized = []
    for img in images:
        if img.height != h:
            w = int(img.width * (h / img.height))
            img = img.resize((w, h), Image.LANCZOS)
        resized.append(img)
    w_total = sum(img.width for img in resized)
    canvas = Image.new("RGB", (w_total, h), bg_color)
    x = 0
    for img in resized:
        canvas.paste(img, (x, 0))
        x += img.width
    return canvas

# NEW: リサイズユーティリティ
def resize_to_preset(img: Image.Image, preset_key: str) -> Image.Image:
    w, h = RESOLUTIONS[preset_key]
    return img.resize((w, h), Image.LANCZOS)

@spaces.GPU()
def generate_turnaround(
    image,
    extra_prompt="",
    preset_key="nealy 9:16",  # NEW: デフォルト
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    progress=gr.Progress(track_tqdm=True),
):
    print(extra_prompt)
    try:
        upload_image_and_prompt_cpu(image, extra_prompt)
    except Exception as e:
        print("Upload failed:", e) 
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    if image is None:
        return None, None, None, None, None, seed, "エラー: 入力画像をアップロードしてください"

    if isinstance(image, Image.Image):
        input_image = image.convert("RGB")
    else:
        input_image = Image.open(image).convert("RGB")

    pil_images = [input_image]

    # 各プロンプト末尾に追記
    p_front = _append_prompt(PROMPTS["front"], extra_prompt)
    p_back  = _append_prompt(PROMPTS["back"],  extra_prompt)
    p_left  = _append_prompt(PROMPTS["left"],  extra_prompt)
    p_right = _append_prompt(PROMPTS["right"], extra_prompt)

    progress(0.25, desc="正面生成中...")
    front = generate_single_view(pil_images, p_front, seed, num_inference_steps, true_guidance_scale)

    progress(0.5, desc="背面生成中...")
    back = generate_single_view([front], p_back, seed+1, num_inference_steps, true_guidance_scale)

    progress(0.75, desc="左側面生成中...")
    left = generate_single_view([front], p_left, seed+2, num_inference_steps, true_guidance_scale)

    progress(1.0, desc="右側面生成中...")
    right = generate_single_view([front], p_right, seed+3, num_inference_steps, true_guidance_scale)

    # NEW: ここで指定プリセットにリサイズ
    front_r = resize_to_preset(front, preset_key)
    back_r  = resize_to_preset(back,  preset_key)
    left_r  = resize_to_preset(left,  preset_key)
    right_r = resize_to_preset(right, preset_key)

    # NEW: リサイズ後を連結（横：正面→右→背面→左）
    concat = concat_images_horizontally([front_r, right_r, back_r, left_r])

    return front_r, back_r, left_r, right_r, concat, seed, f"✅ {preset_key} にリサイズして4視点＋連結画像を生成しました"

# --- UI ---
css = """
#col-container {margin: 0 auto; max-width: 1400px;}
.image-container img {object-fit: contain !important; max-width: 100%; max-height: 100%;}
/* 追加: 注意ボックスのスタイル */
.notice {
  background: #fff5f5;
  border: 1px solid #fca5a5;
  color: #7f1d1d;
  padding: 12px 14px;
  border-radius: 10px;
  font-weight: 600;
  line-height: 1.5;
  margin-bottom: 10px;
}
"""

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

        input_image = gr.Image(label="入力画像", type="pil", height=500)

        # 追記プロンプト欄
        extra_prompt = gr.Textbox(
            label="追加プロンプト（各視点プロンプトの末尾に追記）",
            placeholder="例: high detail, anime style, soft lighting, 4k, pastel colors",
            lines=2
        )

        # NEW: 出力解像度プリセットのプルダウン
        preset_dropdown = gr.Dropdown(
            label="出力解像度プリセット",
            choices=list(RESOLUTIONS.keys()),
            value="nealy 9:16"
        )

        run_button = gr.Button("🎨 生成開始", variant="primary")
        status_text = gr.Textbox(label="ステータス", interactive=False)

        with gr.Row():
            result_front = gr.Image(label="正面", type="pil", format="png", height=400, show_download_button=True)
            result_back = gr.Image(label="背面", type="pil", format="png", height=400, show_download_button=True)
        with gr.Row():
            result_left = gr.Image(label="左側面", type="pil", format="png", height=400, show_download_button=True)
            result_right = gr.Image(label="右側面", type="pil", format="png", height=400, show_download_button=True)

        # PNG連結出力
        result_concat = gr.Image(label="連結画像（正面→右→背面→左）", type="pil", format="png", height=400, show_download_button=True)

        with gr.Accordion("⚙️ 詳細設定", open=False):
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            randomize_seed = gr.Checkbox(label="ランダムシード", value=True)
            true_guidance_scale = gr.Slider(label="True guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
            num_inference_steps = gr.Slider(label="生成ステップ数", minimum=1, maximum=40, step=1, value=4)

    # NEW: クリック時に preset_dropdown を引数として渡す
    run_button.click(
        fn=generate_turnaround,
        inputs=[input_image, extra_prompt, preset_dropdown, seed, randomize_seed, true_guidance_scale, num_inference_steps],
        outputs=[result_front, result_back, result_left, result_right, result_concat, seed, status_text],
    )

if __name__ == "__main__":
    demo.launch()