update

app.py

@@ -1,7 +1,166 @@
 import gradio as gr
+from ominicontrol import generate_image
 
-def greet(name):
-    return "Hello " + name + "!!"
+import spaces
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+USE_ZERO_GPU = True
+
+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;
+}
+"""
+
+
+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("# 🌍 OminiControl (Image Stylization)")
+        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")
+
+            with gr.Column(elem_classes="outputPanel"):
+                output_images = gr.Image(
+                    type="pil",
+                    width=640,
+                    height=640,
+                    label="Output Image",
+                )
+
+        btn.click(
+            fn=infer,
+            inputs=[
+                style,
+                original_image,
+                inference_mode,
+                image_guidance,
+                image_ratio,
+                use_random_seed,
+                seed,
+                steps,
+            ],
+            outputs=output_images,
+        )
+
+    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 = 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
+
+
+if USE_ZERO_GPU:
+    infer = spaces.GPU(infer, duration=360)
+
+if __name__ == "__main__":
+    demo = gradio_interface()
+    demo.launch(
+        debug=True,
+        server_name="0.0.0.0",
+    )

ominicontrol.py

@@ -0,0 +1,129 @@
+import torch
+from diffusers.pipelines import FluxPipeline
+from OminiControl.src.flux.condition import Condition
+from PIL import Image
+import random
+import os
+
+from OminiControl.src.flux.generate import generate, seed_everything
+
+HF_TOKEN=os.getenv("HF_TOKEN")
+
+print("Loading model...")
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, use_auth_token=HF_TOKEN
+)
+pipe = pipe.to("cuda")
+
+pipe.unload_lora_weights()
+
+pipe.load_lora_weights(
+    "Yuanshi/OminiControlStyle",
+    weight_name=f"v0/ghibli.safetensors",
+    adapter_name="ghibli",
+    use_auth_token=HF_TOKEN
+)
+pipe.load_lora_weights(
+    "Yuanshi/OminiControlStyle",
+    weight_name=f"v0/irasutoya.safetensors",
+    adapter_name="irasutoya",
+    use_auth_token=HF_TOKEN
+)
+pipe.load_lora_weights(
+    "Yuanshi/OminiControlStyle",
+    weight_name=f"v0/simpsons.safetensors",
+    adapter_name="simpsons",
+    use_auth_token=HF_TOKEN
+)
+pipe.load_lora_weights(
+    "Yuanshi/OminiControlStyle",
+    weight_name=f"v0/snoopy.safetensors",
+    adapter_name="snoopy",
+    use_auth_token=HF_TOKEN
+)
+
+
+
+def generate_image(
+    image,
+    style,
+    inference_mode,
+    image_guidance,
+    image_ratio,
+    steps,
+    use_random_seed,
+    seed,
+):
+    # 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]
+
+    return result_img