Update app.py

app.py

@@ -11,37 +11,53 @@ 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
-
 
+# --- Model & Repo ---
 HF_MODEL = os.environ.get("HF_UPLOAD_REPO", "rahul7star/qwen-edit-img-repo")
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# --- Camera prompts ---
+BASE_PROMPTS = {
+    "front": "Move the camera to a front-facing position showing the full character. Background is plain white.",
+    "back": "Move the camera to a back-facing position showing the full character. Background is plain white.",
+    "left": "Move the camera to a side (left) profile view. Background is plain white.",
+    "right": "Move the camera to a side (right) profile view. Background is plain white.",
+    "45_left": "Rotate camera 45° left",
+    "45_right": "Rotate camera 45° right",
+    "90_left": "Rotate camera 90° left",
+    "90_right": "Rotate camera 90° right",
+    "top_down": "Switch to top-down view",
+    "low_angle": "Switch to low-angle view",
+    "close_up": "Switch to close-up lens",
+    "medium_close_up": "Switch to medium close-up lens",
+    "zoom_out": "Switch to zoom out lens",
+}
+
+# --- Resolution presets ---
+RESOLUTIONS = {
+    "1:4": (512, 2048),
+    "1:3": (576, 1728),
+    "nealy 9:16": (768, 1344),
+    "nealy 2:3": (832, 1216),
+    "3:4": (896, 1152),
+}
+
+MAX_SEED = np.iinfo(np.int32).max
 
 # --- CPU-only upload function ---
 def upload_image_and_prompt_cpu(input_image, prompt_text) -> str:
     from datetime import datetime
-    import tempfile, os, uuid, shutil
+    import 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)
@@ -49,7 +65,6 @@ def upload_image_and_prompt_cpu(input_image, prompt_text) -> str:
             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",
@@ -58,7 +73,6 @@ def upload_image_and_prompt_cpu(input_image, prompt_text) -> str:
         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)
@@ -71,17 +85,11 @@ def upload_image_and_prompt_cpu(input_image, prompt_text) -> str:
         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 & model load ---
 scheduler_config = {
     "base_image_seq_len": 256,
     "base_shift": math.log(3),
@@ -98,66 +106,38 @@ scheduler_config = {
     "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)
 
+# Load LoRA weights
 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),
-}
-
+# --- Utilities ---
 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,
@@ -173,6 +153,10 @@ def generate_single_view(input_images, prompt, seed, num_inference_steps, true_g
         print("Upload failed:", e)
     return result[0]
 
+def resize_to_preset(img: Image.Image, preset_key: str) -> Image.Image:
+    w, h = RESOLUTIONS[preset_key]
+    return img.resize((w, h), Image.LANCZOS)
+
 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:
@@ -192,131 +176,84 @@ def concat_images_horizontally(images, bg_color=(255, 255, 255)):
         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)
-
+# --- Main generation function ---
 @spaces.GPU()
 def generate_turnaround(
     image,
+    selected_angles,
     extra_prompt="",
-    preset_key="nealy 9:16",  # NEW: デフォルト
+    preset_key="nealy 9:16",
     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)
+        return {}, seed, "❌ 入力画像をアップロードしてください"
 
-    progress(0.5, desc="背面生成中...")
-    back = generate_single_view([front], p_back, seed+1, num_inference_steps, true_guidance_scale)
+    input_image = image.convert("RGB") if isinstance(image, Image.Image) else Image.open(image).convert("RGB")
 
-    progress(0.75, desc="左側面生成中...")
-    left = generate_single_view([front], p_left, seed+2, num_inference_steps, true_guidance_scale)
+    results = {}
+    current_seed = seed
+    for i, angle in enumerate(selected_angles):
+        progress((i+1)/len(selected_angles), desc=f"{angle} 生成中...")
+        prompt = _append_prompt(BASE_PROMPTS[angle], extra_prompt)
+        img = generate_single_view([input_image], prompt, current_seed, num_inference_steps, true_guidance_scale)
+        img = resize_to_preset(img, preset_key)
+        results[angle] = img
+        current_seed += 1
 
-    progress(1.0, desc="右側面生成中...")
-    right = generate_single_view([front], p_right, seed+3, num_inference_steps, true_guidance_scale)
+    # Concatenate all selected images in order
+    concat_img = concat_images_horizontally(list(results.values()))
+    results["concat"] = concat_img
 
-    # 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視点＋連結画像を生成しました"
+    return results, seed, f"✅ {preset_key} にリサイズして {len(selected_angles)} 視点＋連結画像を生成しました"
 
 # --- 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
+        extra_prompt = gr.Textbox(label="追加プロンプト", placeholder="high detail, anime style, soft lighting, 4k", lines=2)
+        preset_dropdown = gr.Dropdown(label="出力解像度プリセット", choices=list(RESOLUTIONS.keys()), value="nealy 9:16")
+        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+        randomize_seed = gr.Checkbox(label="ランダムシード", value=True)
+
+        # --- Checklist for angles ---
+        select_all_checkbox = gr.Checkbox(label="Select All", value=True)
+        angles_checklist = gr.CheckboxGroup(
+            label="生成するカメラ視点を選択",
+            choices=list(BASE_PROMPTS.keys()),
+            value=list(BASE_PROMPTS.keys())
         )
 
-        # NEW: 出力解像度プリセットのプルダウン
-        preset_dropdown = gr.Dropdown(
-            label="出力解像度プリセット",
-            choices=list(RESOLUTIONS.keys()),
-            value="nealy 9:16"
+        # JS: update checklist when select_all changes
+        select_all_checkbox.change(
+            lambda select_all: list(BASE_PROMPTS.keys()) if select_all else [],
+            inputs=select_all_checkbox,
+            outputs=angles_checklist
         )
 
         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)
+        # Dynamic image outputs
+        image_outputs = {angle: gr.Image(label=angle, type="pil", format="png", height=400, show_download_button=True)
+                         for angle in BASE_PROMPTS.keys()}
+        image_outputs["concat"] = gr.Image(label="連結画像", type="pil", format="png", height=400, show_download_button=True)
 
-    # NEW: クリック時に preset_dropdown を引数として渡す
+    # Button click
     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],
+        inputs=[input_image, angles_checklist, extra_prompt, preset_dropdown, seed, randomize_seed, seed, 4],
+        outputs=[image_outputs, seed, status_text],
     )
 
 if __name__ == "__main__":