Upload 6 files

.gitattributes

@@ -1,35 +1,53 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/cartoon_boy.png filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/clock.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/demo/demo_this_is_omini_control.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/demo/dreambooth_res.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/demo/monalisa_omini.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/demo/scene_variation.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/demo/try_on.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/monalisa.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/rc_car.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/room_corner.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/tshirt.jpg filter=lfs diff=lfs merge=lfs -text
+OminiControl/assets/vase_hq.jpg filter=lfs diff=lfs merge=lfs -text
+examples/breakingbad.jpg filter=lfs diff=lfs merge=lfs -text
+examples/DistractedBoyfriend.webp filter=lfs diff=lfs merge=lfs -text
+examples/doge.jpg filter=lfs diff=lfs merge=lfs -text
+examples/oiiai.png filter=lfs diff=lfs merge=lfs -text
+examples/PulpFiction.jpg filter=lfs diff=lfs merge=lfs -text
+examples/steve.webp filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,13 @@
----
-title: Ghibli Art
-emoji: 📊
-colorFrom: indigo
-colorTo: yellow
-sdk: gradio
-sdk_version: 5.25.2
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: OminiControl Art
+emoji: 🎨
+colorFrom: green
+colorTo: red
+sdk: gradio
+sdk_version: 5.23.2
+app_file: app.py
+pinned: false
+license: unknown
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,240 @@
+import gradio as gr
+import spaces
+from ominicontrol import generate_image, vote_feedback
+import os
+
+
+USE_ZERO_GPU = os.environ.get("USE_ZERO_GPU", "0") == "1"
+
+css = """
+.inputPanel {
+    width: 320px;
+    display: flex;
+    align-items: center;
+}
+.outputPanel {
+    display: flex;
+    align-items: center;
+}
+.hint {
+    font-size: 14px;
+    color: #777;
+    # border: 1px solid #ccc;
+    padding: 4px;
+    border-radius: 5px;
+    # background-color: #efefef;
+}
+"""
+
+header = """
+# 🎨 OminiControl Art
+<div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
+<a href="https://arxiv.org/abs/2411.15098"><img src="https://img.shields.io/badge/ariXv-Paper-A42C25.svg" alt="arXiv"></a>
+<a href="https://huggingface.co/spaces/Yuanshi/OminiControl"><img src="https://img.shields.io/badge/🤗OminiControl-Demo-ffbd45.svg" alt="HuggingFace"></a>
+<a href="https://github.com/Yuanshi9815/OminiControl"><img src="https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github&" alt="GitHub"></a>
+</div>
+
+***OminiControl Art*** distills the artistic style of [GPT-4o](https://openai.com/index/introducing-4o-image-generation/) into the [FLUX.1](https://blackforestlabs.ai/) model, building on the foundation of [OminiControl](https://github.com/Yuanshi9815/OminiControl)✨.
+Enjoy playing around! 🌈
+"""
+
+
+def style_transfer(image, style):
+    return image
+
+
+styles = [
+    "Studio Ghibli",
+    "Irasutoya Illustration",
+    "The Simpsons",
+    "Snoopy",
+]
+
+
+def gradio_interface():
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown(header)
+        with gr.Row(equal_height=False):
+            with gr.Column(variant="panel", elem_classes="inputPanel"):
+                original_image = gr.Image(
+                    type="pil",
+                    label="Condition Image",
+                    width=400,
+                    height=400,
+                )
+                style = gr.Radio(
+                    styles,
+                    label="🎨 Select Style",
+                    value=styles[0],
+                )
+                # Advanced settings
+                with gr.Accordion(
+                    "⚙️ Advanced Settings", open=False
+                ) as advanced_settings:
+                    inference_mode = gr.Radio(
+                        ["High Quality", "Fast"],
+                        value="High Quality",
+                        label="Generating Mode",
+                    )
+                    image_ratio = gr.Radio(
+                        ["Auto", "Square(1:1)", "Portrait(2:3)", "Landscape(3:2)"],
+                        label="Image Ratio",
+                        value="Auto",
+                    )
+                    use_random_seed = gr.Checkbox(label="Use Random Seed", value=True)
+                    seed = gr.Number(
+                        label="Seed",
+                        value=42,
+                        visible=(not use_random_seed.value),
+                    )
+                    use_random_seed.change(
+                        lambda x: gr.update(visible=(not x)),
+                        use_random_seed,
+                        seed,
+                        show_progress="hidden",
+                    )
+                    image_guidance = gr.Slider(
+                        label="Image Guidance",
+                        minimum=1.1,
+                        maximum=5,
+                        value=1.5,
+                        step=0.1,
+                    )
+                    steps = gr.Slider(
+                        label="Steps",
+                        minimum=10,
+                        maximum=50,
+                        value=20,
+                        step=1,
+                    )
+                    inference_mode.change(
+                        lambda x: gr.update(interactive=(x == "High Quality")),
+                        inference_mode,
+                        image_guidance,
+                        show_progress="hidden",
+                    )
+
+                btn = gr.Button("Generate Image", variant="primary")
+
+                with gr.Accordion("🏞️ Examples", open=True) as advanced_settings:
+                    examples = gr.Examples(
+                        examples=[
+                            ["examples/DistractedBoyfriend.webp", styles[0]],
+                            ["examples/steve.webp", styles[0]],
+                            ["examples/oiiai.png", styles[1]],
+                            ["examples/doge.jpg", styles[1]],
+                            ["examples/breakingbad.jpg", styles[2]],
+                            ["examples/PulpFiction.jpg", styles[3]],
+                        ],
+                        inputs=[original_image, style],
+                    )
+
+            with gr.Column(elem_classes="outputPanel"):
+                output_image = gr.Image(
+                    type="pil",
+                    width=600,
+                    height=600,
+                    label="Output Image",
+                    interactive=False,
+                    sources=None,
+                )
+                inference_id = gr.Textbox(
+                    visible=False,
+                    interactive=False,
+                )
+
+                # Feedback buttons
+                with gr.Column(visible=False) as feedback:
+                    gr.Markdown(
+                        """
+                        Your feedback improves the model! Please let us know how you feel about the generated image.
+                        """,
+                    )
+                    with gr.Row() as feedback_buttons:
+                        upvote = gr.Button("👍 I like it", variant="primary")
+                        downvote = gr.Button("👎 It looks bad")
+
+                        def feedback_func(feedback):
+                            def func(inputs):
+                                print(f"Feedback: {feedback}, Inference ID: {inputs}")
+                                vote_feedback(log_id=inputs, feedback=feedback)
+                                # Here you can add your feedback logging logic
+                                return gr.update(visible=False)
+
+                            return func
+
+                        upvote.click(feedback_func("1"), inference_id, feedback)
+                        downvote.click(feedback_func("0"), inference_id, feedback)
+
+                inference_id.change(
+                    lambda x: gr.update(visible=True), output_image, feedback
+                )
+
+                hint = gr.Markdown(
+                    """
+                    <div style="text-align: center; width: 100%;">
+                    <b>Note: The selected style is in beta testing.</b> Feel free to try a few more times to get a the better result.
+                    </div>
+                    """,
+                    visible=False,
+                )
+                style.change(
+                    lambda x: gr.update(visible=x in styles[1:]),
+                    style,
+                    hint,
+                )
+
+        # with gr.Row():
+        btn.click(
+            fn=infer,
+            inputs=[
+                style,
+                original_image,
+                inference_mode,
+                image_guidance,
+                image_ratio,
+                use_random_seed,
+                seed,
+                steps,
+            ],
+            outputs=[
+                output_image,
+                inference_id,
+            ],
+        )
+
+    return demo
+
+
+def infer(
+    style,
+    original_image,
+    inference_mode,
+    image_guidance,
+    image_ratio,
+    use_random_seed,
+    seed,
+    steps,
+):
+    print(
+        f"Style: {style}, Inference Mode: {inference_mode}, Image Guidance: {image_guidance}, Image Ratio: {image_ratio}, Use Random Seed: {use_random_seed}, Seed: {seed}"
+    )
+    result_image, inference_id = generate_image(
+        image=original_image,
+        style=style,
+        inference_mode=inference_mode,
+        image_guidance=image_guidance,
+        image_ratio=image_ratio,
+        use_random_seed=use_random_seed,
+        seed=seed,
+        steps=steps,
+    )
+    return result_image, inference_id
+
+
+if USE_ZERO_GPU:
+    infer = spaces.GPU(infer)
+
+if __name__ == "__main__":
+    demo = gradio_interface()
+    demo.launch(server_name="0.0.0.0", ssr_mode=False)

log.py

@@ -0,0 +1,78 @@
+import boto3
+import uuid
+import time
+import os
+
+from PIL import Image
+from io import BytesIO
+
+
+MAX_PIXELS = 2048
+
+AWS_BUCKET_NAME = os.environ.get("AWS_BUCKET_NAME", "")
+AWS_INFERENCE_LOG_TABLE = os.environ.get("AWS_INFERENCE_LOG_TABLE", "")
+AWS_FEEDBACK_LOG_TABLE = os.environ.get("AWS_FEEDBACK_LOG_TABLE", "")
+
+
+AWS_REGION = os.environ.get("AWS_REGION", "")
+AWS_ACCESS_ID = os.environ.get("AWS_ACCESS_ID", "")
+AWS_ACCESS_KEY = os.environ.get("AWS_ACCESS_KEY", "")
+
+
+aws_cfg = {
+    "aws_access_key_id": AWS_ACCESS_ID,
+    "aws_secret_access_key": AWS_ACCESS_KEY,
+    "region_name": AWS_REGION,
+}
+
+s3_client = boto3.client("s3", **aws_cfg)
+dynamodb = boto3.resource("dynamodb", **aws_cfg)
+
+inference_log = dynamodb.Table(AWS_INFERENCE_LOG_TABLE)
+feedback_log = dynamodb.Table(AWS_FEEDBACK_LOG_TABLE)
+
+
+def get_metadata():
+    return {
+        "_id": uuid.uuid4().hex,
+        "created_at": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+    }
+
+
+def insert_log(table_type: str, data: dict):
+    assert table_type in ["inference", "feedback"], "Invalid table type"
+    table = inference_log if table_type == "inference" else feedback_log
+    metadata = get_metadata()
+    response = table.put_item(
+        Item={
+            **data,
+            **metadata,
+        }
+    )
+    return response, metadata["_id"]
+
+
+# Example usage:
+# insert_log("inference", {"data": "test"})
+# insert_log("feedback", {"data": "test"})
+
+
+def get_image_obj(image: Image) -> BytesIO:
+    image.thumbnail((MAX_PIXELS, MAX_PIXELS))
+    image_obj = BytesIO()
+    image.save(image_obj, format="WEBP")
+    image_obj.seek(0)
+    return image_obj
+
+
+def log_image(image: Image) -> str:
+    metadata = get_metadata()
+    image_obj = get_image_obj(image)
+    s3_key = f"images/{metadata['_id']}.webp"
+    s3_client.upload_fileobj(image_obj, AWS_BUCKET_NAME, s3_key)
+    return metadata["_id"]
+
+
+# Example usage:
+# image = Image.open("examples/doge.jpg")
+# log_image(image)

ominicontrol.py

@@ -0,0 +1,158 @@
+import torch
+from diffusers.pipelines import FluxPipeline
+from OminiControl.src.flux.condition import Condition
+from PIL import Image
+import random
+
+from OminiControl.src.flux.generate import generate, seed_everything
+
+from log import insert_log, log_image
+
+print("Loading model...")
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
+)
+pipe = pipe.to("cuda")
+
+pipe.unload_lora_weights()
+
+pipe.load_lora_weights(
+    "Yuanshi/OminiControlStyle",
+    weight_name=f"v0/ghibli.safetensors",
+    adapter_name="ghibli",
+)
+pipe.load_lora_weights(
+    "Yuanshi/OminiControlStyle",
+    weight_name=f"v0/irasutoya.safetensors",
+    adapter_name="irasutoya",
+)
+pipe.load_lora_weights(
+    "Yuanshi/OminiControlStyle",
+    weight_name=f"v0/simpsons.safetensors",
+    adapter_name="simpsons",
+)
+pipe.load_lora_weights(
+    "Yuanshi/OminiControlStyle",
+    weight_name=f"v0/snoopy.safetensors",
+    adapter_name="snoopy",
+)
+
+
+def generate_image(
+    image,
+    style,
+    inference_mode,
+    image_guidance,
+    image_ratio,
+    steps,
+    use_random_seed,
+    seed,
+):  
+    condition_id = log_image(image)
+
+    # Prepare Condition
+    def resize(img, factor=16):
+        w, h = img.size
+        new_w, new_h = w // factor * factor, h // factor * factor
+        padding_w, padding_h = (w - new_w) // 2, (h - new_h) // 2
+        img = img.crop((padding_w, padding_h, new_w + padding_w, new_h + padding_h))
+        return img
+
+    # Set Adapter
+    activate_adapter_name = {
+        "Studio Ghibli": "ghibli",
+        "Irasutoya Illustration": "irasutoya",
+        "The Simpsons": "simpsons",
+        "Snoopy": "snoopy",
+    }[style]
+    pipe.set_adapters(activate_adapter_name)
+
+    factor = 512 / max(image.size)
+    image = resize(
+        image.resize(
+            (int(image.size[0] * factor), int(image.size[1] * factor)),
+            Image.LANCZOS,
+        )
+    )
+    delta = -image.size[0] // 16
+    condition = Condition(
+        "subject",
+        # activate_adapter_name,
+        image,
+        position_delta=(0, delta),
+    )
+
+    # Prepare seed
+    if use_random_seed:
+        seed = random.randint(0, 2**32 - 1)
+    seed_everything(seed)
+
+    # Image guidance scale
+    image_guidance = 1.0 if inference_mode == "Fast" else image_guidance
+
+    # Output size
+    if image_ratio == "Auto":
+        r = image.size[0] / image.size[1]
+        ratio = min([0.67, 1, 1.5], key=lambda x: abs(x - r))
+    else:
+        ratio = {
+            "Square(1:1)": 1,
+            "Portrait(2:3)": 0.67,
+            "Landscape(3:2)": 1.5,
+        }[image_ratio]
+    width, height = {
+        0.67: (640, 960),
+        1: (640, 640),
+        1.5: (960, 640),
+    }[ratio]
+
+    print(
+        f"Image Ratio: {image_ratio}, Inference Mode: {inference_mode}, Image Guidance: {image_guidance}, Seed: {seed}, Steps: {steps}, Size: {width}x{height}"
+    )
+    # Generate
+    result_img = generate(
+        pipe,
+        prompt="",
+        conditions=[condition],
+        num_inference_steps=steps,
+        width=width,
+        height=height,
+        image_guidance_scale=image_guidance,
+        default_lora=True,
+        max_sequence_length=32,
+    ).images[0]
+    # result_img = image
+    result_id = log_image(result_img)
+
+    log_data = {
+        "condition": condition_id,
+        "result": result_id,
+        "prompt": "",
+        "inference_mode": inference_mode,
+        "image_guidance_scale": image_guidance,
+        "seed": seed,
+        "steps": steps,
+        "style": style,
+        "width": width,
+        "height": height,
+    }
+    log_data = {k: str(v) for k, v in log_data.items()}
+
+    _, log_id = insert_log("inference", log_data)
+
+    print(f"Image log ID: {log_id}")
+
+    return result_img, log_id
+
+
+def vote_feedback(
+    log_id,
+    feedback,
+):
+    log_data = {
+        "log_id": log_id,
+        "feedback": feedback,
+    }
+    log_data = {k: str(v) for k, v in log_data.items()}
+
+    insert_log("feedback", log_data)

requirements.txt

@@ -0,0 +1,10 @@
+transformers
+diffusers
+peft
+opencv-python
+protobuf
+sentencepiece
+jupyter
+torchao
+
+boto3