mcp + cleanup

.gitignore

@@ -0,0 +1,2 @@
+**/__pycache__
+**/.DS_Store

app.py

@@ -46,25 +46,59 @@ def load_pipeline(accelerator, weight_dtype, args):
 
 @spaces.GPU(duration=180)
 def run(
-    instruction,
-    width_input,
-    height_input,
-    scheduler,
-    num_inference_steps,
-    image_input_1,
-    image_input_2,
-    image_input_3,
-    negative_prompt,
-    guidance_scale_input,
-    img_guidance_scale_input,
-    cfg_range_start,
-    cfg_range_end,
-    num_images_per_prompt,
-    max_input_image_side_length,
-    max_pixels,
-    seed_input,
+    prompt,
+    width_input=1024,
+    height_input=1024,
+    scheduler='euler',
+    num_inference_steps=50,
+    image_input_1=None,
+    image_input_2=None,
+    image_input_3=None,
+    negative_prompt=NEGATIVE_PROMPT,
+    guidance_scale_input=5.0,
+    img_guidance_scale_input=2.0,
+    cfg_range_start=0.0,
+    cfg_range_end=1.0,
+    num_images_per_prompt=1,
+    max_input_image_side_length=2048,
+    max_pixels=1024 * 1024,
+    seed_input=-1,
     progress=gr.Progress(),
 ):
+    """
+    Generate and/or edit images based on text prompts and optional input images using the OmniGen2 pipeline.
+    
+    Quick Tips for Best Results:
+    - Generate high-resolution images (at least 512x512 recommended).
+    - Be Specific: Instead of "Add bird to desk", try "Add the bird from image 1 to the desk in image 2".
+    - Use English: English prompts currently yield better results.
+    - Adjust image_guidance_scale for better consistency with the reference image:
+        - Image Editing: 1.3 - 2.0
+        - In-context Generation: 2.0 - 3.0
+    
+    Args:
+        prompt: Text prompt describing the desired image generation or editing task.
+        width_input: Width of the output image in pixels (default: 1024).
+        height_input: Height of the output image in pixels (default: 1024).
+        scheduler: Scheduler type to use ('euler' or 'dpmsolver') (default: 'euler').
+        num_inference_steps: Number of denoising steps (default: 50).
+        image_input_1: Optional first input image for editing/reference (default: None).
+        image_input_2: Optional second input image for editing/reference (default: None).
+        image_input_3: Optional third input image for editing/reference (default: None).
+        negative_prompt: Text describing what to avoid in the generation (default: NEGATIVE_PROMPT).
+        guidance_scale_input: Text guidance scale for controlling prompt adherence (default: 5.0).
+        img_guidance_scale_input: Image guidance scale for controlling input image adherence (default: 2.0).
+        cfg_range_start: Start of the classifier-free guidance range (0.0-1.0) (default: 0.0).
+        cfg_range_end: End of the classifier-free guidance range (0.0-1.0) (default: 1.0).
+        num_images_per_prompt: Number of images to generate per prompt (default: 1).
+        max_input_image_side_length: Maximum side length for input images (default: 2048).
+        max_pixels: Maximum total pixels for input images (default: 1048576).
+        seed_input: Random seed for reproducible generation (-1 for random) (default: -1).
+        progress: Gradio progress tracker. (Can be None when not ran on Gradio)
+        
+    Returns:
+        PIL.Image: Generated image or collage of multiple generated images.
+    """
     input_images = [image_input_1, image_input_2, image_input_3]
     input_images = [img for img in input_images if img is not None]
 
@@ -91,7 +125,7 @@ def run(
         )
 
     results = pipeline(
-        prompt=instruction,
+        prompt=prompt,
         input_images=input_images,
         width=width_input,
         height=height_input,
@@ -135,480 +169,6 @@ def run(
     return output_image
 
 
-def get_example():
-    case = [
-        [
-            "The sun rises slightly, the dew on the rose petals in the garden is clear, a crystal ladybug is crawling to the dew, the background is the early morning garden, macro lens.",
-            1024,
-            1024,
-            'euler',
-            50,
-            None,
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            3.5,
-            1.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "A snow maiden with pale translucent skin, frosty white lashes, and a soft expression of longing",
-            1024,
-            1024,
-            'euler',
-            50,
-            None,
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            3.5,
-            1.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Add a fisherman hat to the woman's head",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/flux5.png"),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            " replace the sword with a hammer.",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(
-                ROOT_DIR,
-                "example_images/d8f8f44c64106e7715c61b5dfa9d9ca0974314c5d4a4a50418acf7ff373432bb.png",
-            ),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Extract the character from the picture and fill the rest of the background with white.",
-            # "Transform the sculpture into jade",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(
-                ROOT_DIR, "example_images/46e79704-c88e-4e68-97b4-b4c40cd29826.png"
-            ),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Make he smile",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(
-                ROOT_DIR, "example_images/vicky-hladynets-C8Ta0gwPbQg-unsplash.jpg"
-            ),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Change the background to classroom",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/ComfyUI_temp_mllvz_00071_.png"),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Raise his hand",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(
-                ROOT_DIR,
-                "example_images/289089159-a6d7abc142419e63cab0a566eb38e0fb6acb217b340f054c6172139b316f6596.png",
-            ),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Generate a photo of an anime-style figurine placed on a desk. The figurine model should be based on the character photo provided in the attachment, accurately replicating the full-body pose, facial expression, and clothing style of the character in the photo, ensuring the entire figurine is fully presented. The overall design should be exquisite and detailed, soft gradient colors and a delicate texture, leaning towards a Japanese anime style, rich in details, with a realistic quality and beautiful visual appeal.",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/RAL_0315.JPG"),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Change the dress to blue.",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/1.png"),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Remove the cat",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(
-                ROOT_DIR,
-                "example_images/386724677-589d19050d4ea0603aee6831459aede29a24f4d8668c62c049f413db31508a54.png",
-            ),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "In a cozy café, the anime figure is sitting in front of a laptop, smiling confidently.",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/ComfyUI_00254_.png"),
-            None,
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Create a wedding figure based on the girl in the first image and the man in the second image. Set the background as a wedding hall, with the man dressed in a suit and the girl in a white wedding dress. Ensure that the original faces remain unchanged and are accurately preserved. The man should adopt a realistic style, whereas the girl should maintain their classic anime style.",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/1_20241127203215.png"),
-            os.path.join(ROOT_DIR, "example_images/000050281.jpg"),
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            3.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Let the girl  and the boy get married in the church. ",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/8FtFUxRzXqaguVRGzkHvN.png"),
-            os.path.join(ROOT_DIR, "example_images/01194-20240127001056_1024x1536.png"),
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            3.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Let the man from image1 and the woman from image2 kiss and hug",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/1280X1280.png"),
-            os.path.join(ROOT_DIR, "example_images/000077066.jpg"),
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Please let the person in image 2 hold the toy from the first image in a parking lot.",
-            1024,
-            1024,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/04.jpg"),
-            os.path.join(ROOT_DIR, "example_images/000365954.jpg"),
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Make the girl pray in the second image.",
-            1024,
-            682,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/000440817.jpg"),
-            os.path.join(ROOT_DIR, "example_images/000119733.jpg"),
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Add the bird from image 1 to the desk in image 2",
-            1024,
-            682,
-            'euler',
-            50,
-            os.path.join(
-                ROOT_DIR,
-                "example_images/996e2cf6-daa5-48c4-9ad7-0719af640c17_1748848108409.png",
-            ),
-            os.path.join(ROOT_DIR, "example_images/00066-10350085.png"),
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Replace the apple in the first image with the cat from the second image",
-            1024,
-            780,
-            'euler',
-            50,
-            os.path.join(ROOT_DIR, "example_images/apple.png"),
-            os.path.join(
-                ROOT_DIR,
-                "example_images/468404374-d52ec1a44aa7e0dc9c2807ce09d303a111c78f34da3da2401b83ce10815ff872.png",
-            ),
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "Replace the woman in the second image with the woman from the first image",
-            1024,
-            747,
-            'euler',
-            50,
-            os.path.join(
-                ROOT_DIR, "example_images/byward-outfitters-B97YFrsITyo-unsplash.jpg"
-            ),
-            os.path.join(
-                ROOT_DIR, "example_images/6652baf6-4096-40ef-a475-425e4c072daf.png"
-            ),
-            None,
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-        [
-            "The cat is sitting on the table. The bird is perching on the edge of the table.",
-            800,
-            512,
-            "euler",
-            50,
-            os.path.join(
-                ROOT_DIR,
-                "example_images/996e2cf6-daa5-48c4-9ad7-0719af640c17_1748848108409.png",
-            ),
-            os.path.join(
-                ROOT_DIR,
-                "example_images/468404374-d52ec1a44aa7e0dc9c2807ce09d303a111c78f34da3da2401b83ce10815ff872.png",
-            ),
-            os.path.join(ROOT_DIR, "example_images/00066-10350085.png"),
-            NEGATIVE_PROMPT,
-            5.0,
-            2.0,
-            0.0,
-            1.0,
-            1,
-            2048,
-            1024 * 1024,
-            0,
-        ],
-    ]
-    return case
-
-
-def run_for_examples(
-    instruction,
-    width_input,
-    height_input,
-    scheduler,
-    num_inference_steps,
-    image_input_1,
-    image_input_2,
-    image_input_3,
-    negative_prompt,
-    text_guidance_scale_input,
-    image_guidance_scale_input,
-    cfg_range_start,
-    cfg_range_end,
-    num_images_per_prompt,
-    max_input_image_side_length,
-    max_pixels,
-    seed_input,
-):
-    return run(
-        instruction,
-        width_input,
-        height_input,
-        scheduler,
-        num_inference_steps,
-        image_input_1,
-        image_input_2,
-        image_input_3,
-        negative_prompt,
-        text_guidance_scale_input,
-        image_guidance_scale_input,
-        cfg_range_start,
-        cfg_range_end,
-        num_images_per_prompt,
-        max_input_image_side_length,
-        max_pixels,
-        seed_input,
-    )
-
 description = """
 ### 💡 Quick Tips for Best Results (see our [github](https://github.com/VectorSpaceLab/OmniGen2?tab=readme-ov-file#-usage-tips) for more details)
 - Image Quality: Use high-resolution images (at least 512x512 recommended).
@@ -640,8 +200,8 @@ def main(args):
         with gr.Row():
             with gr.Column():
                 # text prompt
-                instruction = gr.Textbox(
-                    label='Enter your prompt. Use "first/second image" or “第一张图/第二张图” as reference.',
+                prompt = gr.Textbox(
+                    label='Enter your prompt. Use "first/second image" as reference.',
                     placeholder="Type your prompt here...",
                 )
 
@@ -777,32 +337,7 @@ def main(args):
         generate_button.click(
             run,
             inputs=[
-                instruction,
-                width_input,
-                height_input,
-                scheduler_input,
-                num_inference_steps,
-                image_input_1,
-                image_input_2,
-                image_input_3,
-                negative_prompt,
-                text_guidance_scale_input,
-                image_guidance_scale_input,
-                cfg_range_start,
-                cfg_range_end,
-                num_images_per_prompt,
-                max_input_image_side_length,
-                max_pixels,
-                seed_input,
-            ],
-            outputs=output_image,
-        )
-
-        gr.Examples(
-            examples=get_example(),
-            fn=run_for_examples,
-            inputs=[
-                instruction,
+                prompt,
                 width_input,
                 height_input,
                 scheduler_input,
@@ -825,7 +360,7 @@ def main(args):
 
         gr.Markdown(article)
     # launch
-    demo.launch(share=args.share, server_port=args.port, allowed_paths=[ROOT_DIR])
+    demo.launch(share=args.share, server_port=args.port, allowed_paths=[ROOT_DIR], mcp_server=True)
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Run the OmniGen2")