app.py

import gradio as gr
import os
import random
from PIL import Image
import spaces
import torch
from transformers import MllamaForConditionalGeneration, AutoProcessor
from OmniGen import OmniGenPipeline

Llama32V_HFtoken = os.getenv("Llama32V")
pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1")
model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
model = MllamaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto",hf_token = Llama32V_HFtoken)
processor = AutoProcessor.from_pretrained(model_id)

@spaces.GPU()
def predict_clothing(images):
    messages = [{"role": "user", "content":
                 [
                     {"type": "image"},
                     {"type": "text", "text": "Define this clothing in 1-3 words. Your response should be only the definition."}
                 ]}
               ]
    input_text = processor.apply_chat_template(messages, add_generation_prompt=True)

    output_texts = []
    for image in images:
        inputs = processor(image, input_text, add_special_tokens=False, return_tensors="pt").to(model.device)
        with torch.no_grad():
            output = model.generate(**inputs, max_new_tokens=30)
        output_texts.append(str(processor.decode(output[0])))

    return output_texts
    

@spaces.GPU(duration=180)
def generate_image(img1, img2, img3, height, width, img_guidance_scale, inference_steps, seed, separate_cfg_infer, offload_model,
                   use_input_image_size_as_output, max_input_image_size, randomize_seed, guidance_scale=3.5):
    
    input_images = [img1, img2, img3]
    # Delete None
    input_images = [img for img in input_images if img is not None]
    if len(input_images) == 0:
        input_images = None

    wears = predict_clothing(input_images[1:])
    if len(wears)==1:
        dress = wears[0]
        text = """A male wearing a {dress}. The male is in <img><|image_1|></img>. The {dress} is in <img><|image_2|></img>."""
    elif len(wears)==2:
        topwear, bottomwear = wears[0], wears[1]
        text = """A male wearing a {topwear} and a {bottomwear}. The male is in <img><|image_1|></img>.
                    The {topwear} is in <img><|image_2|></img>. The {bottomwear} is in <img><|image_3|></img>."""
    else:
        input_images = None
        
    
    if randomize_seed:
        seed = random.randint(0, 10000000)

    output = pipe(prompt=text, input_images=input_images, height=height, width=width, guidance_scale=guidance_scale,
                  img_guidance_scale=img_guidance_scale, num_inference_steps=inference_steps, separate_cfg_infer=separate_cfg_infer, 
                  use_kv_cache=True, offload_kv_cache=True, offload_model=offload_model, 
                  use_input_image_size_as_output=use_input_image_size_as_output, seed=seed, max_input_image_size=max_input_image_size,)
    img = output[0]
    return img

def get_example():
    case = [
        [
            "A curly-haired man in a red shirt is drinking tea.",
            None,
            None,
            None,
            1024,
            1024,
            2.5,
            1.6,
            0,
            1024,
            False,
            False,
        ],
        [
            "The woman in <img><|image_1|></img> waves her hand happily in the crowd",
            "./imgs/test_cases/zhang.png",
            None,
            None,
            1024,
            1024,
            2.5,
            1.9,
            128,
            1024,
            False,
            False,
        ],
        [
            "A man in a black shirt is reading a book. The man is the right man in <img><|image_1|></img>.",
            "./imgs/test_cases/two_man.jpg",
            None,
            None,
            1024,
            1024,
            2.5,
            1.6,
            0,
            1024,
            False,
            False,
        ],
        [
            "Two woman are raising fried chicken legs in a bar. A woman is <img><|image_1|></img>. Another woman is <img><|image_2|></img>.",
            "./imgs/test_cases/mckenna.jpg",
            "./imgs/test_cases/Amanda.jpg",
            None,
            1024,
            1024,
            2.5,
            1.8,
            65,
            1024,
            False,
            False,
        ],
        [
            "A man and a short-haired woman with a wrinkled face are standing in front of a bookshelf in a library. The man is the man in the middle of <img><|image_1|></img>, and the woman is oldest woman in <img><|image_2|></img>",
            "./imgs/test_cases/1.jpg",
            "./imgs/test_cases/2.jpg",
            None,
            1024,
            1024,
            2.5,
            1.6,
            60,
            1024,
            False,
            False,
        ],
        [
            "A man and a woman are sitting at a classroom desk. The man is the man with yellow hair in <img><|image_1|></img>. The woman is the woman on the left of <img><|image_2|></img>",
            "./imgs/test_cases/3.jpg",
            "./imgs/test_cases/4.jpg",
            None,
            1024,
            1024,
            2.5,
            1.8,
            66,
            1024,
            False,
            False,
        ],
        [
            "The flower <img><|image_1|></img> is placed in the vase which is in the middle of <img><|image_2|></img> on a wooden table of a living room",
            "./imgs/test_cases/rose.jpg",
            "./imgs/test_cases/vase.jpg",
            None,
            1024,
            1024,
            2.5,
            1.6,
            0,
            1024,
            False,
            False,
        ],
        [
            "<img><|image_1|><img>\n Remove the woman's earrings. Replace the mug with a clear glass filled with sparkling iced cola.",
            "./imgs/demo_cases/t2i_woman_with_book.png",
            None,
            None,
            None,
            None,
            2.5,
            1.6,
            222,
            1024,
            False,
            True,
        ],
        [
            "Detect the skeleton of human in this image: <img><|image_1|></img>.",
            "./imgs/test_cases/control.jpg",
            None,
            None,
            1024,
            1024,
            2.0,
            1.6,
            0,
            1024,
            False,
            True,
        ],
        [
            "Generate a new photo using the following picture and text as conditions: <img><|image_1|><img>\n A young boy is sitting on a sofa in the library, holding a book. His hair is neatly combed, and a faint smile plays on his lips, with a few freckles scattered across his cheeks. The library is quiet, with rows of shelves filled with books stretching out behind him.",
            "./imgs/demo_cases/skeletal.png",
            None,
            None,
            1024,
            1024,
            2,
            1.6,
            999,
            1024,
            False,
            True,
        ],
        [
            "Following the pose of this image <img><|image_1|><img>, generate a new photo: A young boy is sitting on a sofa in the library, holding a book. His hair is neatly combed, and a faint smile plays on his lips, with a few freckles scattered across his cheeks. The library is quiet, with rows of shelves filled with books stretching out behind him.",
            "./imgs/demo_cases/edit.png",
            None,
            None,
            1024,
            1024,
            2.0,
            1.6,
            123,
            1024,
            False,
            True,
        ],
        [
            "Following the depth mapping of this image <img><|image_1|><img>, generate a new photo: A young girl is sitting on a sofa in the library, holding a book. His hair is neatly combed, and a faint smile plays on his lips, with a few freckles scattered across his cheeks. The library is quiet, with rows of shelves filled with books stretching out behind him.",
            "./imgs/demo_cases/edit.png",
            None,
            None,
            1024,
            1024,
            2.0,
            1.6,
            1,
            1024,
            False,
            True,
        ],
        [
            "<img><|image_1|><\/img> What item can be used to see the current time? Please highlight it in blue.",
            "./imgs/test_cases/watch.jpg",
            None,
            None,
            1024,
            1024,
            2.5,
            1.6,
            666,
            1024,
            False,
            True,
        ],
        [
            "According to the following examples, generate an output for the input.\nInput: <img><|image_1|></img>\nOutput: <img><|image_2|></img>\n\nInput: <img><|image_3|></img>\nOutput: ",
            "./imgs/test_cases/icl1.jpg",
            "./imgs/test_cases/icl2.jpg",
            "./imgs/test_cases/icl3.jpg",
            224,
            224,
            2.5,
            1.6,
            1,
            768,
            False,
            False,
        ],
    ]
    return case

def run_for_examples(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, seed, max_input_image_size, randomize_seed, use_input_image_size_as_output):    
    # Check the internal configuration of the function
    inference_steps = 50
    separate_cfg_infer = True
    offload_model = False
    
    return generate_image(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed,
                          separate_cfg_infer, offload_model, use_input_image_size_as_output, max_input_image_size, randomize_seed)

description = """
This is a Virtual Try-On Platform.
Usage:
- First upload your own image as the first image, also tagged 'Person'
- Then upload you 'Top-wear' and 'Bottom-wear' images
- If its a single dress, and/or you don't have a Topwear and Bottomwear as separate images upload that single image under 'Topwear'

Tips:
- For image editing task and controlnet task, we recommend setting the height and width of output image as the same as input image. For example, if you want to edit a 512x512 image, you should set the height and width of output image as 512x512. You also can set the `use_input_image_size_as_output` to automatically set the height and width of output image as the same as input image.
- For out-of-memory or time cost, you can set `offload_model=True` or refer to [./docs/inference.md#requiremented-resources](https://github.com/VectorSpaceLab/OmniGen/blob/main/docs/inference.md#requiremented-resources) to select a appropriate setting.
- If inference time is too long when inputting multiple images, please try to reduce the `max_input_image_size`. For more details please refer to [./docs/inference.md#requiremented-resources](https://github.com/VectorSpaceLab/OmniGen/blob/main/docs/inference.md#requiremented-resources).

**HF Spaces often encounter errors due to quota limitations, so recommend to run it locally.**
"""

Credits = """**Credits** 
Made using [OmniGen](https://huggingface.co/Shitao/OmniGen-v1): Unified Image Generation [paper](https://arxiv.org/abs/2409.11340) [code](https://github.com/VectorSpaceLab/OmniGen)
"""


# Gradio 
with gr.Blocks() as demo:
    gr.Markdown("Virtual Try-On")
    gr.Markdown(description)
    with gr.Row():
        with gr.Row(equal_height=True):
            # input images
            image_input_1 = gr.Image(label="Person", type="filepath")
            image_input_2 = gr.Image(label="Top-wear", type="filepath")
            image_input_3 = gr.Image(label="Bottom-wear", type="filepath")

        # slider
        height_input = gr.Slider(
            label="Height", minimum=128, maximum=2048, value=1024, step=16
        )
        width_input = gr.Slider(
            label="Width", minimum=128, maximum=2048, value=1024, step=16
        )

        guidance_scale_input = gr.Slider(
            label="Guidance Scale", minimum=1.0, maximum=5.0, value=2.5, step=0.1
        )

        img_guidance_scale_input = gr.Slider(
            label="img_guidance_scale", minimum=1.0, maximum=2.0, value=1.6, step=0.1
        )

        num_inference_steps = gr.Slider(
            label="Inference Steps", minimum=1, maximum=100, value=50, step=1
        )

        seed_input = gr.Slider(
            label="Seed", minimum=0, maximum=2147483647, value=42, step=1
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        max_input_image_size = gr.Slider(
            label="max_input_image_size", minimum=128, maximum=2048, value=1024, step=16
        )

        separate_cfg_infer = gr.Checkbox(
            label="separate_cfg_infer", info="Whether to use separate inference process for different guidance. This will reduce the memory cost.", value=True,
        )
        offload_model = gr.Checkbox(
            label="offload_model", info="Offload model to CPU, which will significantly reduce the memory cost but slow down the generation speed. You can cancel separate_cfg_infer and set offload_model=True. If both separate_cfg_infer and offload_model are True, further reduce the memory, but slowest generation", value=False,
        )
        use_input_image_size_as_output = gr.Checkbox(
            label="use_input_image_size_as_output", info="Automatically adjust the output image size to be same as input image size. For editing and controlnet task, it can make sure the output image has the same size as input image leading to better performance", value=False,
        )

        # generate
        generate_button = gr.Button("Generate Image")
            
    with gr.Row():
        # output image
        output_image = gr.Image(label="Output Image")

    # click
    generate_button.click(
        generate_image,
        inputs=[image_input_1, image_input_2, image_input_3, height_input, width_input, img_guidance_scale_input, num_inference_steps,
                seed_input, separate_cfg_infer, offload_model, use_input_image_size_as_output, max_input_image_size, randomize_seed,
                guidance_scale_input,],
        outputs=output_image,
    )

    gr.Examples(
        examples=get_example(),
        fn=run_for_examples,
        inputs=[image_input_1, image_input_2, image_input_3, height_input, width_input, img_guidance_scale_input, seed_input,
                max_input_image_size, randomize_seed, use_input_image_size_as_output,guidance_scale_input],
        outputs=output_image,
    )

    gr.Markdown(Credits)

# launch
demo.launch()