Add application file

.gitignore

@@ -0,0 +1,4 @@
+gradio_cached_examples/
+flagged/
+.idea/
+.DS_Store

app.py

@@ -0,0 +1,133 @@
+import os
+import shlex
+import subprocess
+
+if os.getenv('SYSTEM', "") == 'spaces' and not os.getenv('USE_PRIVATE_PACKAGE', False):
+    GITHUB_TOKEN = os.getenv('GITHUB_TOKEN')
+    GITHUB_USER = os.getenv('GITHUB_USER')
+    git_repo = f"https://{GITHUB_TOKEN}@github.com/{GITHUB_USER}/VIBE.git"
+    subprocess.call(shlex.split(f'pip install git+{git_repo}'))
+
+from functools import partial
+from gradio.components import Image, Textbox
+import gradio as gr
+
+from PIL import Image
+from huggingface_hub import snapshot_download
+import os
+import random
+import torch
+import numpy as np
+import pathlib
+
+from vibe.editor import ImageEditor
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+def load_pipeline():
+    HF_TOKEN = os.getenv('HF_TOKEN')
+    model_path = snapshot_download(
+        repo_id="iitolstykh/VIBE-Image-Edit",
+        repo_type="model",
+        token=HF_TOKEN,
+    )
+
+    # Load model
+    editor_pipeline = ImageEditor(
+        checkpoint_path=model_path,
+        image_guidance_scale=1.2,
+        guidance_scale=4.5,
+        num_inference_steps=20,
+        device="cuda:0",
+    )
+
+    return editor_pipeline
+
+def set_env(seed=0):
+    torch.manual_seed(seed)
+    torch.set_grad_enabled(False)
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+@torch.inference_mode()
+def generate_img(
+    pipeline: ImageEditor,
+    np_image: np.ndarray,  # rgb image
+    edit_prompt: str,
+    sample_steps,
+    scale,
+    image_guidance_scale,
+    seed=0,
+    randomize_seed=False,
+):
+    # color_converted_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
+
+    pil_image = Image.fromarray(np_image)
+
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    set_env(seed)
+
+    edited_image = pipeline.generate_edited_image(
+        instruction=edit_prompt,
+        conditioning_image=pil_image,
+        num_images_per_prompt=1,
+        num_inference_steps=sample_steps,
+        guidance_scale=scale,
+        image_guidance_scale=image_guidance_scale,
+    )[0]
+    return edited_image
+
+
+if __name__ == "__main__":
+
+    DESCRIPTION = f"""DEMO for VIBE-Image-Edit model."""
+
+    image_dir = pathlib.Path('images')
+    examples = [[path.as_posix(), "let this case swim in the river", 20, 4.5, 1.2, 42] for path in sorted(image_dir.glob('*.png'))]
+
+    editor_pipeline = load_pipeline()
+    generate_fn = partial(
+        generate_img,
+        editor_pipeline=editor_pipeline
+    )
+
+    demo = gr.Interface(
+        fn=generate_fn,
+        inputs=[
+            gr.Image(label="Input", type="numpy"),
+            Textbox(label="Prompt", placeholder="Please enter your prompt. \n"),
+            gr.Slider(label="Sample Steps", minimum=1, maximum=100, value=20, step=1),
+            gr.Slider(
+                label="Guidance Scale", minimum=0.1, maximum=30.0, value=4.5, step=0.1
+            ),
+            gr.Slider(
+                label="Image Guidance Scale",
+                minimum=0.1,
+                maximum=30.0,
+                value=1.2,
+                step=0.1,
+            ),
+            gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=42,
+            ),
+            gr.Checkbox(label="Randomize seed", value=False),
+        ],
+        outputs=[
+            Image(type="numpy", label="Img"),
+        ],
+        title="",
+        description=DESCRIPTION,
+        examples=examples,
+    )
+
+    demo.queue(max_size=15).launch()

requirements.txt

@@ -0,0 +1,8 @@
+transformers==4.57.1
+torchvision==0.21.0
+torch==2.6.0
+diffusers==0.33.1
+loguru==0.7.3
+gradio==6.2.0
+huggingface_hub==0.34.1
+# git+https://github.com/ai-forever/VIBE