Rename gradio.py to inference.py

gradio.py → inference.py

@@ -1,282 +1,260 @@
-import gradio as gr
-import numpy as np
+# inference.py
+import os
+import argparse
 import random
-import torch
-import spaces
+import json
+import base64
+from io import BytesIO
+from datetime import datetime
 
+import torch
+import numpy as np
+import requests
 from PIL import Image
 from diffusers import QwenImageEditPipeline
 
-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.  
+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**. Avoid overly long sentences and reduce unnecessary descriptive language.  
-- If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary.  
-- Keep the core intention 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 edited input image’s overall scene.  
+- Keep the rewritten prompt **concise**. Avoid overly long sentences and reduce unnecessary descriptive language.  
+- If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary.  
+- Keep the core intention 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 edited input image’s overall scene.  
 
 ## 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.  
+- 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 `" "`. Do not translate or alter the original language of the text, and do not change the capitalization.  
+- All text content must be enclosed in English double quotes `" "`. Do not translate or alter the original language of the text, and do not change the capitalization.  
 - **For text replacement tasks, always use the fixed template:**
-    - `Replace "xx" to "yy"`.  
-    - `Replace the xx bounding box to "yy"`.  
-- If the user does not specify text content, infer and add concise text based on the instruction and the input image’s context. For example:  
-    > Original: "Add a line of text" (poster)  
-    > Rewritten: "Add text \"LIMITED EDITION\" at the top center with slight shadow"  
-- Specify text position, color, and layout in a concise way.  
+    - `Replace "xx" to "yy"`.  
+    - `Replace the xx bounding box to "yy"`.  
+- If the user does not specify text content, infer and add concise text based on the instruction and the input image’s context. For example:  
+    > Original: "Add a line of text" (poster)  
+    > Rewritten: "Add text \"LIMITED EDITION\" at the top center with slight shadow"  
+- Specify text position, color, and layout in a concise way.  
 
 ### 3. Human Editing Tasks
-- Maintain the person’s core visual consistency (ethnicity, gender, age, hairstyle, expression, outfit, etc.).  
-- If modifying appearance (e.g., clothes, hairstyle), ensure the new element is consistent with the original style.  
-- **For expression changes, they must be natural and subtle, never exaggerated.**  
+- Maintain the person’s core visual consistency (ethnicity, gender, age, hairstyle, expression, outfit, etc.).  
+- If modifying appearance (e.g., clothes, hairstyle), ensure the new element is consistent with the original style.  
+- **For expression changes, they must be natural and subtle, never exaggerated.**  
 - If deletion is not specifically emphasized, the most important subject in the original image (e.g., a person, an animal) should be preserved.
-    - For background change tasks, emphasize maintaining subject consistency at first.  
-- Example:  
-    > Original: "Change the person’s hat"  
-    > Rewritten: "Replace the man’s hat with a dark brown beret; keep smile, short hair, and gray jacket unchanged"  
+    - For background change tasks, emphasize maintaining subject consistency at first.  
+- Example:  
+    > Original: "Change the person’s hat"  
+    > Rewritten: "Replace the man’s hat with a dark brown beret; keep smile, short hair, and gray jacket unchanged"  
 
 ### 4. Style Transformation or Enhancement Tasks
-- If a style is specified, describe it concisely with key visual traits. For example:  
-    > Original: "Disco style"  
-    > Rewritten: "1970s disco: flashing lights, disco ball, mirrored walls, colorful tones"  
-- If the instruction says "use reference style" or "keep current style," analyze the input image, extract main features (color, composition, texture, lighting, art style), and integrate them concisely.  
-- **For coloring tasks, including restoring old photos, always use the fixed template:** "Restore old photograph, remove scratches, reduce noise, enhance details, high resolution, realistic, natural skin tones, clear facial features, no distortion, vintage photo restoration"  
+- If a style is specified, describe it concisely with key visual traits. For example:  
+    > Original: "Disco style"  
+    > Rewritten: "1970s disco: flashing lights, disco ball, mirrored walls, colorful tones"  
+- If the instruction says "use reference style" or "keep current style," analyze the input image, extract main features (color, composition, texture, lighting, art style), and integrate them concisely.  
+- **For coloring tasks, including restoring old photos, always use the fixed template:** "Restore old photograph, remove scratches, reduce noise, enhance details, high resolution, realistic, natural skin tones, clear facial features, no distortion, vintage photo restoration"  
 - If there are other changes, place the style description at the end.
 
 ## 3. Rationality and Logic Checks
-- Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" should be logically corrected.  
-- Add missing key information: if position is unspecified, choose a reasonable area based on composition (near subject, empty space, center/edges).  
+- Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" should be logically corrected.  
+- Add missing key information: if position is unspecified, choose a reasonable area based on composition (near subject, empty space, center/edges).  
 
 # Output Format Example
 ```json
 {
-   "Rewritten": "..."
+   "Rewritten": "..."
 }
 '''
 
 def polish_prompt(prompt, img):
-    prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
-    success=False
-    while not success:
+    """使用 DashScope API 重写和优化提示词"""
+    if not os.environ.get('DASH_API_KEY'):
+        print("[警告] 环境变量 DASH_API_KEY 未设置，将跳过提示词重写。")
+        return prompt
+
+    full_prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
+    for attempt in range(3): # 最多重试3次
         try:
-            result = api(prompt, [img])
-            # print(f"Result: {result}")
-            # print(f"Polished Prompt: {polished_prompt}")
+            result = api(full_prompt, [img])
             if isinstance(result, str):
-                result = result.replace('```json','')
-                result = result.replace('```','')
-                result = json.loads(result)
+                result_json_str = result.replace('```json','').replace('```','')
+                result_data = json.loads(result_json_str)
             else:
-                result = json.loads(result)
-
-            polished_prompt = result['Rewritten']
-            polished_prompt = polished_prompt.strip()
-            polished_prompt = polished_prompt.replace("\n", " ")
-            success = True
+                result_data = json.loads(result)
+            
+            polished = result_data['Rewritten']
+            return polished.strip().replace("\n", " ")
         except Exception as e:
-            print(f"[Warning] Error during API call: {e}")
-    return polished_prompt
-
+            print(f"[警告] API调用失败 (尝试 {attempt + 1}): {e}")
+    
+    print("[错误] 多次尝试后提示词重写失败，将使用原始提示词。")
+    return prompt
 
 def encode_image(pil_image):
-    import io
-    buffered = io.BytesIO()
+    """将 PIL 图片编码为 base64 字符串"""
+    buffered = BytesIO()
     pil_image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-
-
-
 def api(prompt, img_list, model="qwen-vl-max-latest", kwargs={}):
+    """调用 DashScope API"""
     import dashscope
     api_key = os.environ.get('DASH_API_KEY')
     if not api_key:
         raise EnvironmentError("DASH_API_KEY is not set")
-    assert model in ["qwen-vl-max-latest"], f"Not implemented model {model}"
-    sys_promot = "you are a helpful assistant, you should provide useful answers to users."
+
     messages = [
-        {"role": "system", "content": sys_promot},
-        {"role": "user", "content": []}]
+        {"role": "system", "content": "you are a helpful assistant, you should provide useful answers to users."},
+        {"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({"image": f"data:image/png;base64,{encode_image(img)}"})
     messages[1]["content"].append({"text": f"{prompt}"})
 
-    response_format = kwargs.get('response_format', None)
-
     response = dashscope.MultiModalConversation.call(
         api_key=api_key,
-        model=model, # For example, use qwen-plus here. You can change the model name as needed. Model list: https://help.aliyun.com/zh/model-studio/getting-started/models
+        model=model,
         messages=messages,
         result_format='message',
-        response_format=response_format,
-        )
+        response_format=kwargs.get('response_format', None),
+    )
 
     if response.status_code == 200:
         return response.output.choices[0].message.content[0]['text']
     else:
         raise Exception(f'Failed to post: {response}')
 
-# --- Model Loading ---
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load the model pipeline
-pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
-
-# --- UI Constants and Helpers ---
-MAX_SEED = np.iinfo(np.int32).max
-
-# --- Main Inference Function (with hardcoded negative prompt) ---
-@spaces.GPU(duration=300)
-def infer(
-    image,
-    prompt,
-    seed=42,
-    randomize_seed=False,
-    true_guidance_scale=1.0,
-    num_inference_steps=50,
-    rewrite_prompt=True,
-    num_images_per_prompt=1,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """
-    Generates an image using the local Qwen-Image diffusers pipeline.
-    """
-    # Hardcode the negative prompt as requested
-    negative_prompt = " "
+def load_image(image_path):
+    """从本地路径或URL加载图片"""
+    try:
+        if image_path.startswith("http://") or image_path.startswith("https://"):
+            response = requests.get(image_path)
+            response.raise_for_status()
+            image = Image.open(BytesIO(response.content)).convert("RGB")
+        else:
+            image = Image.open(image_path).convert("RGB")
+        return image
+    except Exception as e:
+        print(f"❌ 错误：无法加载图片 '{image_path}'。请检查路径或链接是否正确。")
+        print(f"   详细信息: {e}")
+        return None
+
+# --- 主推理逻辑 ---
+
+def main(args):
+    """执行模型推理的主函数"""
+    output_dir = "output"
+    os.makedirs(output_dir, exist_ok=True)
+
+    dtype = torch.bfloat16
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"使用设备: {device}")
+
+    print("正在加载 Qwen-Image-Edit 模型...")
+    try:
+        pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
+        print("模型加载完成。")
+    except Exception as e:
+        print(f"❌ 错误：模型加载失败。请检查网络连接和依赖项。")
+        print(f"   详细信息: {e}")
+        return
+
+    print(f"正在从 '{args.input_image}' 加载输入图片...")
+    input_image = load_image(args.input_image)
+    if input_image is None:
+        return
+
+    # 设置随机种子
+    seed = random.randint(0, np.iinfo(np.int32).max) if args.random_seed else args.seed
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    # 如果不禁用重写功能，则调用 polish_prompt
+    prompt_to_use = polish_prompt(args.prompt, input_image) if not args.no_rewrite else args.prompt
     
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    if not args.no_rewrite:
+        print(f"重写后的提示词: '{prompt_to_use}'")
+
+    print("-" * 30)
+    print("🚀 开始推理...")
+    print(f"  - 提示词: '{prompt_to_use}'")
+    print(f"  - 随机种子: {seed}")
+    print(f"  - 推理步数: {args.steps}")
+    print(f"  -引导系数 (Guidance Scale): {args.guidance_scale}")
+    print("-" * 30)
+
+    try:
+        images = pipe(
+            image=input_image,
+            prompt=prompt_to_use,
+            negative_prompt=" ", # 固定负向���示词
+            num_inference_steps=args.steps,
+            generator=generator,
+            true_cfg_scale=args.guidance_scale,
+            num_images_per_prompt=1
+        ).images
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_path = os.path.join(output_dir, f"output_{timestamp}_{seed}.png")
+        images[0].save(output_path)
+        print(f"✅ 推理成功！图片已保存至: {output_path}")
+
+    except Exception as e:
+        print(f"❌ 推理过程中发生错误: {e}")
+
+# --- 命令行接口 ---
 
-    # Set up the generator for reproducibility
-    generator = torch.Generator(device=device).manual_seed(seed)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Qwen 图像编辑命令行工具")
     
-    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}")
-    if rewrite_prompt:
-        prompt = polish_prompt(prompt, image)
-        print(f"Rewritten Prompt: {prompt}")
-
-    # Generate the image
-    image = pipe(
-        image,
-        prompt=prompt,
-        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 image, seed
-
-# --- Examples and UI Layout ---
-examples = []
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 1024px;
-}
-#edit_text{margin-top: -62px !important}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.HTML('<img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_edit_logo.png" alt="Qwen-Image Logo" width="400" style="display: block; margin: 0 auto;">')
-        gr.Markdown("[Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. Try on [Qwen Chat](https://chat.qwen.ai/), or [download model](https://huggingface.co/Qwen/Qwen-Image-Edit) to run locally with ComfyUI or diffusers.")
-        with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(label="Input Image", show_label=False, type="pil")
-
-            # result = gr.Image(label="Result", show_label=False, type="pil")
-            result = gr.Gallery(label="Result", show_label=False, type="pil")
-        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
-
-            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.Row():
-
-                true_guidance_scale = gr.Slider(
-                    label="True guidance scale",
-                    minimum=1.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=4.0
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=50,
-                )
-                
-                num_images_per_prompt = gr.Slider(
-                    label="Number of images per prompt",
-                    minimum=1,
-                    maximum=4,
-                    step=1,
-                    value=1,
-                )
-                
-                rewrite_prompt = gr.Checkbox(label="Rewrite prompt", value=True)
-
-        # gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)
-
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            input_image,
-            prompt,
-            seed,
-            randomize_seed,
-            true_guidance_scale,
-            num_inference_steps,
-            rewrite_prompt,
-            num_images_per_prompt,
-        ],
-        outputs=[result, seed],
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="必须：用于编辑图像的指令。"
+    )
+    parser.add_argument(
+        "--input_image",
+        type=str,
+        required=True,
+        help="必须：输入图片的本地路径或URL链接。"
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="用于复现结果的随机种子，默认为 42。"
+    )
+    parser.add_argument(
+        "--random_seed",
+        action="store_true",
+        help="如果设置此项，则使用一个随机种子。"
+    )
+    parser.add_argument(
+        "--steps",
+        type=int,
+        default=50,
+        help="推理步数，默认为 50。"
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=4.0,
+        help="引导系数 (CFG scale)，默认为 4.0。"
+    )
+    parser.add_argument(
+        "--no_rewrite",
+        action="store_true",
+        help="如果设置此项，则禁用提示词重写功能。"
     )
 
-if __name__ == "__main__":
-    demo.launch()
+    args = parser.parse_args()
+    main(args)