Update app.py

app.py

@@ -3,165 +3,14 @@ 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
-
-from huggingface_hub import InferenceClient
 import math
-
 import os
-import base64
-import json
-
-SYSTEM_PROMPT = '''
-# Edit Instruction Rewriter
-You are a professional edit instruction rewriter. Your task is to generate a precise, concise, and visually achievable professional-level edit instruction based on the user-provided instruction and the image to be edited.  
-
-Please strictly follow the rewriting rules below:
-
-## 1. General Principles
-- Keep the rewritten prompt **concise and comprehensive**. Avoid overly long sentences and unnecessary descriptive language.  
-- If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary.  
-- Keep the main part of the original instruction unchanged, only enhancing its clarity, rationality, and visual feasibility.  
-- All added objects or modifications must align with the logic and style of the scene in the input images.  
-- If multiple sub-images are to be generated, describe the content of each sub-image individually.  
-
-## 2. Task-Type Handling Rules
-
-### 1. Add, Delete, Replace Tasks
-- If the instruction is clear (already includes task type, target entity, position, quantity, attributes), preserve the original intent and only refine the grammar.  
-- If the description is vague, supplement with minimal but sufficient details (category, color, size, orientation, position, etc.). For example:  
-    > Original: "Add an animal"  
-    > Rewritten: "Add a light-gray cat in the bottom-right corner, sitting and facing the camera"  
-- Remove meaningless instructions: e.g., "Add 0 objects" should be ignored or flagged as invalid.  
-- For replacement tasks, specify "Replace Y with X" and briefly describe the key visual features of X.  
-
-### 2. Text Editing Tasks
-- All text content must be enclosed in English double quotes `" "`. Keep the original language of the text, and keep the capitalization.  
-- Both adding new text and replacing existing text are text replacement tasks, For example:  
-    - Replace "xx" to "yy"  
-    - Replace the mask / bounding box to "yy"  
-    - Replace the visual object to "yy"  
-- Specify text position, color, and layout only if user has required.  
-- If font is specified, keep the original language of the font.  
-
-### 3. Human Editing Tasks
-- Make the smallest changes to the given user's prompt.  
-- If changes to background, action, expression, camera shot, or ambient lighting are required, please list each modification individually.
-- **Edits to makeup or facial features / expression must be subtle, not exaggerated, and must preserve the subject's identity consistency.**
-    > Original: "Add eyebrows to the face"  
-    > Rewritten: "Slightly thicken the person's eyebrows with little change, look natural."
-
-### 4. Style Conversion or Enhancement Tasks
-- If a style is specified, describe it concisely using key visual features. For example:  
-    > Original: "Disco style"  
-    > Rewritten: "1970s disco style: flashing lights, disco ball, mirrored walls, vibrant colors"  
-- For style reference, analyze the original image and extract key characteristics (color, composition, texture, lighting, artistic style, etc.), integrating them into the instruction.  
-- **Colorization tasks (including old photo restoration) must use the fixed template:**  
-  "Restore and colorize the old photo."  
-- Clearly specify the object to be modified. For example:  
-    > Original: Modify the subject in Picture 1 to match the style of Picture 2.  
-    > Rewritten: Change the girl in Picture 1 to the ink-wash style of Picture 2 — rendered in black-and-white watercolor with soft color transitions.
-
-### 5. Material Replacement
-- Clearly specify the object and the material. For example: "Change the material of the apple to papercut style."
-- For text material replacement, use the fixed template:
-    "Change the material of text "xxxx" to laser style"
-
-### 6. Logo/Pattern Editing
-- Material replacement should preserve the original shape and structure as much as possible. For example:
-   > Original: "Convert to sapphire material"  
-   > Rewritten: "Convert the main subject in the image to sapphire material, preserving similar shape and structure"
-- When migrating logos/patterns to new scenes, ensure shape and structure consistency. For example:
-   > Original: "Migrate the logo in the image to a new scene"  
-   > Rewritten: "Migrate the logo in the image to a new scene, preserving similar shape and structure"
-
-### 7. Multi-Image Tasks
-- Rewritten prompts must clearly point out which image's element is being modified. For example:  
-    > Original: "Replace the subject of picture 1 with the subject of picture 2"  
-    > Rewritten: "Replace the girl of picture 1 with the boy of picture 2, keeping picture 2's background unchanged"  
-- For stylization tasks, describe the reference image's style in the rewritten prompt, while preserving the visual content of the source image.  
-
-## 3. Rationale and Logic Check
-- Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" requires logical correction.
-- Supplement missing critical information: e.g., if position is unspecified, choose a reasonable area based on composition (near subject, blank space, center/edge, etc.).
-
-# Output Format Example
-```json
-{
-   "Rewritten": "..."
-}
-'''
-# --- Prompt Enhancement using Hugging Face InferenceClient ---
-def polish_prompt_hf(prompt, img_list):
-    """
-    Rewrites the prompt using a Hugging Face InferenceClient.
-    """
-    # Ensure HF_TOKEN is set
-    api_key = os.environ.get("HF_TOKEN")
-    if not api_key:
-        print("Warning: HF_TOKEN not set. Falling back to original prompt.")
-        return prompt
-
-    try:
-        # Initialize the client
-        prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
-        client = InferenceClient(
-            provider="cerebras",
-            api_key=api_key,
-        )
-
-        # Format the messages for the chat completions API
-        sys_promot = "you are a helpful assistant, you should provide useful answers to users."
-        messages = [
-            {"role": "system", "content": sys_promot},
-            {"role": "user", "content": []}]
-        for img in img_list:
-            messages[1]["content"].append(
-                {"image": f"data:image/png;base64,{encode_image(img)}"})
-        messages[1]["content"].append({"text": f"{prompt}"})
-
-        # Call the API
-        completion = client.chat.completions.create(
-            model="Qwen/Qwen3-235B-A22B-Instruct-2507",
-            messages=messages,
-        )
-        
-        # Parse the response
-        result = completion.choices[0].message.content
-        
-        # Try to extract JSON if present
-        if '{"Rewritten"' in result:
-            try:
-                # Clean up the response
-                result = result.replace('```json', '').replace('```', '')
-                result_json = json.loads(result)
-                polished_prompt = result_json.get('Rewritten', result)
-            except:
-                polished_prompt = result
-        else:
-            polished_prompt = result
-            
-        polished_prompt = polished_prompt.strip().replace("\n", " ")
-        return polished_prompt
-        
-    except Exception as e:
-        print(f"Error during API call to Hugging Face: {e}")
-        # Fallback to original prompt if enhancement fails
-        return prompt
-    
-
-
-def encode_image(pil_image):
-    import io
-    buffered = io.BytesIO()
-    pil_image.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 # --- Model Loading ---
 dtype = torch.bfloat16
@@ -193,8 +42,8 @@ pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509",
                                                  scheduler=scheduler,
                                                  torch_dtype=dtype).to(device)
 pipe.load_lora_weights(
-        "lightx2v/Qwen-Image-Lightning", 
-        weight_name="Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-8steps-V1.0-bf16.safetensors"
+        "2vXpSwA7/iroiro-lora", 
+        weight_name="qwen_lora/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16_dim1.safetensors"
     )
 pipe.fuse_lora()
 
@@ -208,40 +57,54 @@ optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB",
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max
 
-def use_output_as_input(output_images):
-    """Convert output images to input format for the gallery"""
-    if output_images is None or len(output_images) == 0:
-        return []
-    return output_images
+# 固定プロンプト定義
+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. Background is plain white.",
+    "back": "Move the camera to a back-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. 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. The character stands with both arms extended slightly downward and close to the thighs, keeping the body evenly balanced. 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. The character stands with both arms extended slightly downward and close to the thighs, keeping the body evenly balanced. Background is plain white."
+}
 
-# --- Main Inference Function (with hardcoded negative prompt) ---
+def generate_single_view(input_images, prompt, seed, num_inference_steps, true_guidance_scale):
+    """単一の視点の画像を生成"""
+    negative_prompt = " "
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    print(f"Generating with prompt: '{prompt}'")
+    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")
+    
+    # Generate the image
+    result = pipe(
+        image=input_images if input_images else None,
+        prompt=prompt,
+        height=None,
+        width=None,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        true_cfg_scale=true_guidance_scale,
+        num_images_per_prompt=1,
+    ).images
+    
+    return result[0]
+
+# --- Main Inference Function ---
 @spaces.GPU(duration=300)
-def infer(
+def generate_turnaround(
     images,
-    prompt,
     seed=42,
     randomize_seed=False,
     true_guidance_scale=1.0,
-    num_inference_steps=8,
-    height=None,
-    width=None,
-    rewrite_prompt=True,
-    num_images_per_prompt=1,
+    num_inference_steps=4,
     progress=gr.Progress(track_tqdm=True),
 ):
     """
-    Generates an image using the local Qwen-Image diffusers pipeline.
+    入力画像から4つの視点（正面、背面、左側面、右側面）の立ち絵を生成
     """
-    # Hardcode the negative prompt as requested
-    negative_prompt = " "
-    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-
-    # Set up the generator for reproducibility
-    generator = torch.Generator(device=device).manual_seed(seed)
     
-    # Load input images into PIL Images
+    # 入力画像の読み込み
     pil_images = []
     if images is not None:
         for item in images:
@@ -254,87 +117,78 @@ def infer(
                     pil_images.append(Image.open(item.name).convert("RGB"))
             except Exception:
                 continue
-
-    if height==256 and width==256:
-        height, width = None, None
-    print(f"Calling pipeline with prompt: '{prompt}'")
-    print(f"Negative Prompt: '{negative_prompt}'")
-    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
-    if rewrite_prompt and len(pil_images) > 0:
-        prompt = polish_prompt_hf(prompt, pil_images)
-        print(f"Rewritten Prompt: {prompt}")
     
+    if not pil_images:
+        return [None, None, None, None], seed, "エラー: 入力画像をアップロードしてください"
+    
+    results = []
+    
+    # 1. 正面立ち絵を生成
+    progress(0.25, desc="正面立ち絵を生成中...")
+    front_image = generate_single_view(pil_images, PROMPTS["front"], seed, num_inference_steps, true_guidance_scale)
+    results.append(front_image)
+    
+    # 2. 正面立ち絵を入力として背面を生成
+    progress(0.50, desc="背面立ち絵を生成中...")
+    back_image = generate_single_view([front_image], PROMPTS["back"], seed+1, num_inference_steps, true_guidance_scale)
+    results.append(back_image)
+    
+    # 3. 正面立ち絵を入力として左側面を生成
+    progress(0.75, desc="左側面立ち絵を生成中...")
+    left_image = generate_single_view([front_image], PROMPTS["left"], seed+2, num_inference_steps, true_guidance_scale)
+    results.append(left_image)
+    
+    # 4. 正面立ち絵を入力として右側面を生成
+    progress(1.0, desc="右側面立ち絵を生成中...")
+    right_image = generate_single_view([front_image], PROMPTS["right"], seed+3, num_inference_steps, true_guidance_scale)
+    results.append(right_image)
+    
+    return results, seed, "✅ 4視点の立ち絵生成が完了しました"
 
-    # Generate the image
-    image = 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,
-        true_cfg_scale=true_guidance_scale,
-        num_images_per_prompt=num_images_per_prompt,
-    ).images
-
-    # Return images, seed, and make button visible
-    return image, seed, gr.update(visible=True)
-
-# --- Examples and UI Layout ---
-examples = []
-
+# --- UI Layout ---
 css = """
 #col-container {
     margin: 0 auto;
-    max-width: 1024px;
+    max-width: 1200px;
 }
-#logo-title {
+.view-label {
     text-align: center;
+    font-weight: bold;
+    margin-top: 10px;
 }
-#logo-title img {
-    width: 400px;
-}
-#edit_text{margin-top: -62px !important}
 """
 
 with gr.Blocks(css=css) as demo:
+    gr.Markdown("# キャラクター4視点立ち絵自動生成")
+    gr.Markdown("キャラクター画像をアップロードすると、正面・背面・左側面・右側面の4つの立ち絵を自動生成します")
+    
     with gr.Column(elem_id="col-container"):
-        gr.HTML("""
-        <div id="logo-title">
-            <img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_edit_logo.png" alt="Qwen-Image Edit Logo" width="400" style="display: block; margin: 0 auto;">
-            <h2 style="font-style: italic;color: #5b47d1;margin-top: -27px !important;margin-left: 96px">[Plus] Fast, 8-steps with Lightning LoRA</h2>
-        </div>
-        """)
-        gr.Markdown("""
-        [Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. 
-        This demo uses the new [Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) with the [Qwen-Image-Lightning v2](https://huggingface.co/lightx2v/Qwen-Image-Lightning) LoRA + [AoT compilation & FA3](https://huggingface.co/blog/zerogpu-aoti) for accelerated inference.
-        Try on [Qwen Chat](https://chat.qwen.ai/), or [download model](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) to run locally with ComfyUI or diffusers.
-        """)
+        with gr.Row():
+            input_images = gr.Gallery(
+                label="入力画像（キャラクター画像をアップロード）", 
+                show_label=True, 
+                type="pil", 
+                interactive=True,
+                height=400
+            )
+        
+        run_button = gr.Button("🎨 4視点立ち絵を生成", variant="primary", size="lg")
+        
+        status_text = gr.Textbox(label="ステータス", interactive=False)
+        
         with gr.Row():
             with gr.Column():
-                input_images = gr.Gallery(label="Input Images", 
-                                          show_label=False, 
-                                          type="pil", 
-                                          interactive=True)
-
+                result_front = gr.Image(label="正面", type="pil", height=400)
             with gr.Column():
-                result = gr.Gallery(label="Result", show_label=False, type="pil")
-                # Add this button right after the result gallery - initially hidden
-                use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False)
-
+                result_back = gr.Image(label="背面", type="pil", height=400)
+        
         with gr.Row():
-            prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    placeholder="describe the edit instruction",
-                    container=False,
-            )
-            run_button = gr.Button("Edit!", variant="primary")
-
-        with gr.Accordion("Advanced Settings", open=False):
-            # Negative prompt UI element is removed here
+            with gr.Column():
+                result_left = gr.Image(label="左側面", type="pil", height=400)
+            with gr.Column():
+                result_right = gr.Image(label="右側面", type="pil", height=400)
 
+        with gr.Accordion("⚙️ 詳細設定", open=False):
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -343,10 +197,9 @@ with gr.Blocks(css=css) as demo:
                 value=0,
             )
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            randomize_seed = gr.Checkbox(label="ランダムシード", value=True)
 
             with gr.Row():
-
                 true_guidance_scale = gr.Slider(
                     label="True guidance scale",
                     minimum=1.0,
@@ -356,56 +209,31 @@ with gr.Blocks(css=css) as demo:
                 )
 
                 num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
+                    label="生成ステップ数",
                     minimum=1,
                     maximum=40,
                     step=1,
-                    value=8,
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=2048,
-                    step=8,
-                    value=None,
-                )
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=2048,
-                    step=8,
-                    value=None,
+                    value=4,
                 )
-                
-                
-                rewrite_prompt = gr.Checkbox(label="Rewrite prompt", value=False)
 
-        # gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)
+    def update_results(results, seed, status):
+        """結果を個別の画像コンポーネントに分配"""
+        return results[0], results[1], results[2], results[3], seed, status
 
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
+    run_button.click(
+        fn=generate_turnaround,
         inputs=[
             input_images,
-            prompt,
             seed,
             randomize_seed,
             true_guidance_scale,
             num_inference_steps,
-            height,
-            width,
-            rewrite_prompt,
         ],
-        outputs=[result, seed, use_output_btn],  # Added use_output_btn to outputs
-    )
-
-    # Add the new event handler for the "Use Output as Input" button
-    use_output_btn.click(
-        fn=use_output_as_input,
-        inputs=[result],
-        outputs=[input_images]
+        outputs=[gr.State(), seed, status_text],
+    ).then(
+        fn=update_results,
+        inputs=[gr.State(), seed, status_text],
+        outputs=[result_front, result_back, result_left, result_right, seed, status_text]
     )
 
 if __name__ == "__main__":