Create app.py

app.py

@@ -0,0 +1,585 @@
+import spaces
+import argparse
+import os
+import shutil
+import cv2
+import gradio as gr
+import numpy as np
+import torch
+from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+import huggingface_hub
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from torchvision.transforms.functional import normalize
+
+from dreamo.dreamo_pipeline import DreamOPipeline
+from dreamo.utils import img2tensor, resize_numpy_image_area, tensor2img, resize_numpy_image_long
+from tools import BEN2
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--port', type=int, default=8080)
+parser.add_argument('--no_turbo', action='store_true')
+args = parser.parse_args()
+
+huggingface_hub.login(os.getenv('HF_TOKEN'))
+
+try:
+    shutil.rmtree('gradio_cached_examples')
+except FileNotFoundError:
+    print("cache folder not exist")
+
+class Generator:
+    def __init__(self):
+        device = torch.device('cuda')
+        # preprocessing models
+        # background remove model: BEN2
+        self.bg_rm_model = BEN2.BEN_Base().to(device).eval()
+        hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='models')
+        self.bg_rm_model.loadcheckpoints('models/BEN2_Base.pth')
+        # face crop and align tool: facexlib
+        self.face_helper = FaceRestoreHelper(
+            upscale_factor=1,
+            face_size=512,
+            crop_ratio=(1, 1),
+            det_model='retinaface_resnet50',
+            save_ext='png',
+            device=device,
+        )
+
+        # load dreamo
+        model_root = 'black-forest-labs/FLUX.1-dev'
+        dreamo_pipeline = DreamOPipeline.from_pretrained(model_root, torch_dtype=torch.bfloat16)
+        dreamo_pipeline.load_dreamo_model(device, use_turbo=not args.no_turbo)
+        self.dreamo_pipeline = dreamo_pipeline.to(device)
+
+    @torch.no_grad()
+    def get_align_face(self, img):
+        # the face preprocessing code is same as PuLID
+        self.face_helper.clean_all()
+        image_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        self.face_helper.read_image(image_bgr)
+        self.face_helper.get_face_landmarks_5(only_center_face=True)
+        self.face_helper.align_warp_face()
+        if len(self.face_helper.cropped_faces) == 0:
+            return None
+        align_face = self.face_helper.cropped_faces[0]
+
+        input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0
+        input = input.to(torch.device("cuda"))
+        parsing_out = self.face_helper.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0]
+        parsing_out = parsing_out.argmax(dim=1, keepdim=True)
+        bg_label = [0, 16, 18, 7, 8, 9, 14, 15]
+        bg = sum(parsing_out == i for i in bg_label).bool()
+        white_image = torch.ones_like(input)
+        # only keep the face features
+        face_features_image = torch.where(bg, white_image, input)
+        face_features_image = tensor2img(face_features_image, rgb2bgr=False)
+
+        return face_features_image
+
+
+generator = Generator()
+
+
+@spaces.GPU
+@torch.inference_mode()
+def generate_image(
+    ref_image1,
+    ref_image2,
+    ref_task1,
+    ref_task2,
+    prompt,
+    seed,
+    width=1024,
+    height=1024,
+    ref_res=512,
+    num_steps=12,
+    guidance=3.5,
+    true_cfg=1,
+    cfg_start_step=0,
+    cfg_end_step=0,
+    neg_prompt='',
+    neg_guidance=3.5,
+    first_step_guidance=0,
+):
+    print(prompt)
+    ref_conds = []
+    debug_images = []
+
+    ref_images = [ref_image1, ref_image2]
+    ref_tasks = [ref_task1, ref_task2]
+
+    for idx, (ref_image, ref_task) in enumerate(zip(ref_images, ref_tasks)):
+        if ref_image is not None:
+            if ref_task == "id":
+                ref_image = resize_numpy_image_long(ref_image, 1024)
+                ref_image = generator.get_align_face(ref_image)
+            elif ref_task != "style":
+                ref_image = generator.bg_rm_model.inference(Image.fromarray(ref_image))
+            if ref_task != "id":
+                ref_image = resize_numpy_image_area(np.array(ref_image), ref_res * ref_res)
+            debug_images.append(ref_image)
+            ref_image = img2tensor(ref_image, bgr2rgb=False).unsqueeze(0) / 255.0
+            ref_image = 2 * ref_image - 1.0
+            ref_conds.append(
+                {
+                    'img': ref_image,
+                    'task': ref_task,
+                    'idx': idx + 1,
+                }
+            )
+
+    seed = int(seed)
+    if seed == -1:
+        seed = torch.Generator(device="cpu").seed()
+
+    image = generator.dreamo_pipeline(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=num_steps,
+        guidance_scale=guidance,
+        ref_conds=ref_conds,
+        generator=torch.Generator(device="cpu").manual_seed(seed),
+        true_cfg_scale=true_cfg,
+        true_cfg_start_step=cfg_start_step,
+        true_cfg_end_step=cfg_end_step,
+        negative_prompt=neg_prompt,
+        neg_guidance_scale=neg_guidance,
+        first_step_guidance_scale=first_step_guidance if first_step_guidance > 0 else guidance,
+    ).images[0]
+
+    return image, debug_images, seed
+
+
+# Custom CSS for pastel theme
+_CUSTOM_CSS_ = """
+:root {
+    --primary-color: #f8c3cd;            /* Sakura pink - primary accent */
+    --secondary-color: #b3e5fc;          /* Pastel blue - secondary accent */
+    --background-color: #f5f5f7;         /* Very light gray background */
+    --card-background: #ffffff;          /* White for cards */
+    --text-color: #424242;               /* Dark gray for text */
+    --accent-color: #ffb6c1;             /* Light pink for accents */
+    --success-color: #c8e6c9;            /* Pastel green for success */
+    --warning-color: #fff9c4;            /* Pastel yellow for warnings */
+    --shadow-color: rgba(0, 0, 0, 0.1);  /* Shadow color */
+    --border-radius: 12px;               /* Rounded corners */
+}
+
+body {
+    background-color: var(--background-color) !important;
+    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
+}
+
+.gradio-container {
+    max-width: 1200px !important;
+    margin: 0 auto !important;
+}
+
+/* Header styling */
+h1 {
+    color: #9c27b0 !important;
+    font-weight: 800 !important;
+    text-shadow: 2px 2px 4px rgba(156, 39, 176, 0.2) !important;
+    letter-spacing: -0.5px !important;
+}
+
+/* Card styling */
+.gr-box, .gr-form, .gr-panel {
+    border-radius: var(--border-radius) !important;
+    box-shadow: 0 8px 16px var(--shadow-color) !important;
+    background-color: var(--card-background) !important;
+    border: none !important;
+    overflow: hidden !important;
+}
+
+/* Button styling */
+button.gr-button {
+    background: linear-gradient(135deg, var(--primary-color), #e1bee7) !important;
+    border-radius: var(--border-radius) !important;
+    color: #4a148c !important;
+    font-weight: 600 !important;
+    border: none !important;
+    padding: 10px 20px !important;
+    transition: all 0.3s ease !important;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1) !important;
+}
+
+button.gr-button:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 6px 10px rgba(0, 0, 0, 0.15) !important;
+    background: linear-gradient(135deg, #e1bee7, var(--primary-color)) !important;
+}
+
+/* Input fields styling */
+input, select, textarea, .gr-input {
+    border-radius: 8px !important;
+    border: 2px solid #e0e0e0 !important;
+    padding: 10px 15px !important;
+    transition: all 0.3s ease !important;
+    background-color: #fafafa !important;
+}
+
+input:focus, select:focus, textarea:focus, .gr-input:focus {
+    border-color: var(--primary-color) !important;
+    box-shadow: 0 0 0 3px rgba(248, 195, 205, 0.3) !important;
+}
+
+/* Slider styling */
+.gr-slider {
+    --slider-color: var(--primary-color) !important;
+}
+
+.gr-slider-handle {
+    background-color: white !important;
+    border: 2px solid var(--primary-color) !important;
+    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1) !important;
+}
+
+/* Dropdown styling */
+.gr-dropdown {
+    background-color: white !important;
+    border: 2px solid #e0e0e0 !important;
+    border-radius: 8px !important;
+}
+
+.gr-dropdown option {
+    padding: 10px !important;
+}
+
+/* Image upload area */
+.gr-image-input {
+    border: 2px dashed #b39ddb !important;
+    border-radius: var(--border-radius) !important;
+    background-color: #f3e5f5 !important;
+    padding: 20px !important;
+    display: flex !important;
+    flex-direction: column !important;
+    align-items: center !important;
+    justify-content: center !important;
+    transition: all 0.3s ease !important;
+}
+
+.gr-image-input:hover {
+    background-color: #ede7f6 !important;
+    border-color: #9575cd !important;
+}
+
+/* Add a nice pattern to the background */
+body::before {
+    content: "" !important;
+    position: fixed !important;
+    top: 0 !important;
+    left: 0 !important;
+    width: 100% !important;
+    height: 100% !important;
+    background: 
+        radial-gradient(circle at 10% 20%, rgba(248, 195, 205, 0.1) 0%, rgba(245, 245, 247, 0) 20%),
+        radial-gradient(circle at 80% 70%, rgba(179, 229, 252, 0.1) 0%, rgba(245, 245, 247, 0) 20%) !important;
+    pointer-events: none !important;
+    z-index: -1 !important;
+}
+
+/* Gallery styling */
+.gr-gallery {
+    grid-gap: 15px !important;
+}
+
+.gr-gallery-item {
+    border-radius: var(--border-radius) !important;
+    overflow: hidden !important;
+    box-shadow: 0 4px 8px var(--shadow-color) !important;
+    transition: transform 0.3s ease !important;
+}
+
+.gr-gallery-item:hover {
+    transform: scale(1.02) !important;
+}
+
+/* Label styling */
+label {
+    font-weight: 600 !important;
+    color: #673ab7 !important;
+    margin-bottom: 5px !important;
+}
+
+/* Improve spacing */
+.gr-padded {
+    padding: 20px !important;
+}
+
+.gr-compact {
+    gap: 15px !important;
+}
+
+.gr-form > *, .gr-box > * {
+    margin-bottom: 16px !important;
+}
+
+/* Update Examples section */
+.gr-examples {
+    background-color: #f3e5f5 !important;
+    border-radius: var(--border-radius) !important;
+    padding: 15px !important;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.05) !important;
+}
+
+.gr-examples h2 {
+    color: #7b1fa2 !important;
+    font-size: 1.5rem !important;
+    margin-bottom: 15px !important;
+}
+
+/* Dropdown menu adjustments */
+.gr-dropdowns {
+    display: flex !important;
+    gap: 10px !important;
+}
+
+/* Status messages */
+.gr-status {
+    border-radius: 8px !important;
+    padding: 10px 15px !important;
+}
+
+.gr-info {
+    background-color: var(--secondary-color) !important;
+}
+
+.gr-warning {
+    background-color: var(--warning-color) !important;
+}
+
+.gr-error {
+    background-color: #ffcdd2 !important;
+}
+"""
+
+_HEADER_ = '''
+<div style="text-align: center; max-width: 850px; margin: 0 auto; padding: 25px 0;">
+    <div style="background: linear-gradient(135deg, #f8c3cd, #e1bee7, #b3e5fc); color: white; padding: 15px; border-radius: 15px; box-shadow: 0 10px 20px rgba(0,0,0,0.1); margin-bottom: 20px;">
+        <h1 style="font-size: 3rem; font-weight: 800; margin: 0; color: white; text-shadow: 2px 2px 4px rgba(0,0,0,0.2);">✨ DreamO Video ✨</h1>
+        <p style="font-size: 1.2rem; margin: 10px 0 0;">Create customized images with advanced AI</p>
+    </div>
+    <div style="background: white; padding: 15px; border-radius: 12px; box-shadow: 0 5px 15px rgba(0,0,0,0.05);">
+        <p style="font-size: 1rem; margin: 0;">Paper: <a href='https://arxiv.org/abs/2504.16915' target='_blank' style="color: #9c27b0; font-weight: 600;">DreamO: A Unified Framework for Image Customization</a> | 
+        Codes: <a href='https://github.com/bytedance/DreamO' target='_blank' style="color: #9c27b0; font-weight: 600;">GitHub</a></p>
+    </div>
+</div>
+
+<div style="background: #fff9c4; padding: 15px; border-radius: 12px; margin-bottom: 20px; border-left: 5px solid #ffd54f; box-shadow: 0 5px 15px rgba(0,0,0,0.05);">
+    <h3 style="margin-top: 0; color: #ff6f00;">🚩 Update Notes:</h3>
+    <ul style="margin-bottom: 0; padding-left: 20px;">
+        <li><b>2025.05.11:</b> We have updated the model to mitigate over-saturation and plastic-face issues. The new version shows consistent improvements over the previous release.</li>
+        <li><b>2025.05.13:</b> 'DreamO Video' Integration version Release</li>
+    </ul>
+</div>
+'''
+
+_CITE_ = r"""
+<div style="background: white; padding: 20px; border-radius: 12px; margin-top: 20px; box-shadow: 0 5px 15px rgba(0,0,0,0.05);">
+    <p style="margin: 0; font-size: 1.1rem;">If DreamO is helpful, please help to ⭐ the <a href='https://discord.gg/openfreeai' target='_blank' style="color: #9c27b0; font-weight: 600;">community</a>. Thanks!</p>
+    <hr style="border: none; height: 1px; background-color: #e0e0e0; margin: 15px 0;">
+    <h4 style="margin: 0 0 10px; color: #7b1fa2;">📧 Contact</h4>
+    <p style="margin: 0;">If you have any questions or feedback, feel free to open a discussion or contact <b>[email protected]</b></p>
+</div>
+"""
+
+def create_demo():
+    with gr.Blocks(css=_CUSTOM_CSS_) as demo:
+        gr.HTML(_HEADER_)
+
+        with gr.Row():
+            with gr.Column(scale=6):
+                with gr.Box(elem_id="input-panel"):
+                    gr.Markdown("### 📸 Reference Images")
+                    with gr.Row():
+                        with gr.Column():
+                            ref_image1 = gr.Image(label="Reference Image 1", type="numpy", height=256, elem_id="ref-image-1")
+                            ref_task1 = gr.Dropdown(choices=["ip", "id", "style"], value="ip", label="Task for Reference Image 1", elem_id="ref-task-1")
+                        
+                        with gr.Column():
+                            ref_image2 = gr.Image(label="Reference Image 2", type="numpy", height=256, elem_id="ref-image-2")
+                            ref_task2 = gr.Dropdown(choices=["ip", "id", "style"], value="ip", label="Task for Reference Image 2", elem_id="ref-task-2")
+                    
+                    gr.Markdown("### ✏️ Generation Parameters")
+                    prompt = gr.Textbox(label="Prompt", value="a person playing guitar in the street", elem_id="prompt-input")
+                    
+                    with gr.Row():
+                        width = gr.Slider(768, 1024, 1024, step=16, label="Width", elem_id="width-slider")
+                        height = gr.Slider(768, 1024, 1024, step=16, label="Height", elem_id="height-slider")
+                    
+                    with gr.Row():
+                        num_steps = gr.Slider(8, 30, 12, step=1, label="Number of Steps", elem_id="steps-slider")
+                        guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="Guidance Scale", elem_id="guidance-slider")
+                    
+                    seed = gr.Textbox(label="Seed (-1 for random)", value="-1", elem_id="seed-input")
+                    
+                    with gr.Accordion("Advanced Options", open=False):
+                        ref_res = gr.Slider(512, 1024, 512, step=16, label="Resolution for Reference Image")
+                        neg_prompt = gr.Textbox(label="Negative Prompt", value="")
+                        neg_guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="Negative Guidance")
+                        
+                        with gr.Row():
+                            true_cfg = gr.Slider(1, 5, 1, step=0.1, label="True CFG")
+                            first_step_guidance = gr.Slider(0, 10, 0, step=0.1, label="First Step Guidance")
+                        
+                        with gr.Row():
+                            cfg_start_step = gr.Slider(0, 30, 0, step=1, label="CFG Start Step")
+                            cfg_end_step = gr.Slider(0, 30, 0, step=1, label="CFG End Step")
+                    
+                    generate_btn = gr.Button("✨ Generate Image", elem_id="generate-btn")
+                    gr.HTML(_CITE_)
+
+            with gr.Column(scale=6):
+                with gr.Box(elem_id="output-panel"):
+                    gr.Markdown("### 🖼️ Generated Result")
+                    output_image = gr.Image(label="Generated Image", elem_id="output-image", format='png')
+                    seed_output = gr.Textbox(label="Used Seed", elem_id="seed-output")
+                    
+                    gr.Markdown("### 🔍 Preprocessing")
+                    debug_image = gr.Gallery(
+                        label="Preprocessing Results (including face crop and background removal)",
+                        elem_id="debug-gallery",
+                    )
+
+        with gr.Box(elem_id="examples-panel"):
+            gr.Markdown("## 📚 Examples")
+            example_inps = [
+                [
+                    'example_inputs/woman1.png',
+                    None,
+                    'ip',
+                    'ip',
+                    'profile shot dark photo of a 25-year-old female with smoke escaping from her mouth, the backlit smoke gives the image an ephemeral quality, natural face, natural eyebrows, natural skin texture, award winning photo, highly detailed face, atmospheric lighting, film grain, monochrome',  # noqa E501
+                    9180879731249039735,
+                ],
+                [
+                    'example_inputs/man1.png',
+                    None,
+                    'ip',
+                    'ip',
+                    'a man sitting on the cloud, playing guitar',
+                    1206523688721442817,
+                ],
+                [
+                    'example_inputs/toy1.png',
+                    None,
+                    'ip',
+                    'ip',
+                    'a purple toy holding a sign saying "DreamO", on the mountain',
+                    10441727852953907380,
+                ],
+                [
+                    'example_inputs/perfume.png',
+                    None,
+                    'ip',
+                    'ip',
+                    'a perfume under spotlight',
+                    116150031980664704,
+                ],
+                [
+                    'example_inputs/hinton.jpeg',
+                    None,
+                    'id',
+                    'ip',
+                    'portrait, Chibi',
+                    5443415087540486371,
+                ],
+                [
+                    'example_inputs/mickey.png',
+                    None,
+                    'style',
+                    'ip',
+                    'generate a same style image. A rooster wearing overalls.',
+                    6245580464677124951,
+                ],
+                [
+                    'example_inputs/mountain.png',
+                    None,
+                    'style',
+                    'ip',
+                    'generate a same style image. A pavilion by the river, and the distant mountains are endless',
+                    5248066378927500767,
+                ],
+                [
+                    'example_inputs/shirt.png',
+                    'example_inputs/skirt.jpeg',
+                    'ip',
+                    'ip',
+                    'A girl is wearing a short-sleeved shirt and a short skirt on the beach.',
+                    9514069256241143615,
+                ],
+                [
+                    'example_inputs/woman2.png',
+                    'example_inputs/dress.png',
+                    'id',
+                    'ip',
+                    'the woman wearing a dress, In the banquet hall',
+                    7698454872441022867,
+                ],
+                [
+                    'example_inputs/dog1.png',
+                    'example_inputs/dog2.png',
+                    'ip',
+                    'ip',
+                    'two dogs in the jungle',
+                    6187006025405083344,
+                ],
+                [
+                    'example_inputs/woman3.png',
+                    'example_inputs/cat.png',
+                    'ip',
+                    'ip',
+                    'A girl rides a giant cat, walking in the noisy modern city. High definition, realistic, non-cartoonish. Excellent photography work, 8k high definition.',  # noqa E501
+                    11980469406460273604,
+                ],
+                [
+                    'example_inputs/man2.jpeg',
+                    'example_inputs/woman4.jpeg',
+                    'ip',
+                    'ip',
+                    'a man is dancing with a woman in the room',
+                    8303780338601106219,
+                ],
+            ]
+            gr.Examples(
+                examples=example_inps,
+                inputs=[ref_image1, ref_image2, ref_task1, ref_task2, prompt, seed],
+                label='Examples by category: IP task (rows 1-4), ID task (row 5), Style task (rows 6-7), Try-On task (rows 8-9), Multi IP (rows 10-12)',
+                cache_examples='lazy',
+                outputs=[output_image, debug_image, seed_output],
+                fn=generate_image,
+            )
+
+        generate_btn.click(
+            fn=generate_image,
+            inputs=[
+                ref_image1,
+                ref_image2,
+                ref_task1,
+                ref_task2,
+                prompt,
+                seed,
+                width,
+                height,
+                ref_res,
+                num_steps,
+                guidance,
+                true_cfg,
+                cfg_start_step,
+                cfg_end_step,
+                neg_prompt,
+                neg_guidance,
+                first_step_guidance,
+            ],
+            outputs=[output_image, debug_image, seed_output],
+        )
+
+    return demo
+
+
+if __name__ == '__main__':
+    demo = create_demo()
+    demo.launch()