添加prompt tab

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-from ominicontrol import generate_image
+from ominicontrol import generate_image, generate_image_with_prompt
 import os
 
 from huggingface_hub import login
@@ -193,10 +193,134 @@ def infer(
     )
     return result_image
 
+def prompt_gradio_interface():
+    with gr.Blocks(css=css) as demo:
+        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,
+                )
+                prompt = gr.Textbox(
+                    label="Prompt",
+                )
+                # 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.Column(elem_classes="outputPanel"):
+                output_image = gr.Image(
+                    type="pil",
+                    width=600,
+                    height=600,
+                    label="Output Image",
+                    interactive=False,
+                    sources=None,
+                )
+
+        # with gr.Row():
+        btn.click(
+            fn=prompt_infer,
+            inputs=[
+                original_image,
+                prompt,
+                inference_mode,
+                image_guidance,
+                image_ratio,
+                use_random_seed,
+                seed,
+                steps,
+            ],
+            outputs=[
+                output_image,
+            ],
+        )
+
+    return demo
+
+def prompt_infer(
+    original_image,
+    prompt,
+    inference_mode,
+    image_guidance,
+    image_ratio,
+    use_random_seed,
+    seed,
+    steps,
+):
+    result_image = generate_image_with_prompt(
+        image=original_image,
+        prompt=prompt,
+        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
+
+def multi_gradio_interface():
+    with gr.Blocks(css="style.css") as demo:
+        with gr.Tabs():
+            with gr.Tab(label="Style"):
+                gradio_interface()
+            with gr.Tab(label="Prompt"):
+                prompt_gradio_interface()
+
+    return demo
 
 if USE_ZERO_GPU:
     infer = spaces.GPU(infer)
+    prompt_infer = spaces.GPU(prompt_infer)
 
 if __name__ == "__main__":
-    demo = gradio_interface()
+    demo = multi_gradio_interface()
     demo.launch(server_name="0.0.0.0", ssr_mode=False)

ominicontrol.py

@@ -12,6 +12,9 @@ pipe = FluxPipeline.from_pretrained(
 )
 pipe = pipe.to("cuda")
 
+prompt_pipe = FluxPipeline.from_pipe(pipe)
+prompt_pipe = prompt_pipe.to("cuda")
+
 pipe.unload_lora_weights()
 
 pipe.load_lora_weights(
@@ -34,19 +37,6 @@ pipe.load_lora_weights(
     weight_name=f"v0/snoopy.safetensors",
     adapter_name="snoopy",
 )
-# ref: https://civitai.com/models/715472/flux-hayao-miyazaki-ghibli
-pipe.load_lora_weights(
-    "./lora",
-    weight_name="MaoMu_Ghibli.safetensors",
-    adapter_name="MaoMu_Ghibli",
-)
-# ref: https://civitai.com/models/824739/flux-3d-animation-style-lora
-pipe.load_lora_weights(
-    "./lora",
-    weight_name="3d_animation.safetensors",
-    adapter_name="3d_animation",
-)
-
 
 def generate_image(
     image,
@@ -72,8 +62,6 @@ def generate_image(
         "Irasutoya Illustration": "irasutoya",
         "The Simpsons": "simpsons",
         "Snoopy": "snoopy",
-        "3D Animation": "3d_animation",
-        "MaoMu Ghibli": "MaoMu_Ghibli",
     }[style]
     pipe.set_adapters(activate_adapter_name)
 
@@ -145,3 +133,89 @@ def generate_image(
 
     return result_img
 
+
+def generate_image_with_prompt(
+    image,
+    prompt,
+    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
+
+    original_width, original_height = image.size
+
+    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]
+
+
+    output_factor = max(width, height) / max(original_width, original_height)
+    width = int(original_width * output_factor)
+    height = int(original_height * output_factor)
+
+    print(
+        f"Image Ratio: {image_ratio}, Inference Mode: {inference_mode}, Image Guidance: {image_guidance}, Seed: {seed}, Steps: {steps}, Ratio: {ratio}, Size: {width}x{height}"
+    )
+    # Generate
+    result_img = generate(
+        prompt_pipe,
+        prompt=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_img = result_img.resize((width, height), Image.LANCZOS)
+
+    return result_img