prashant-AI-ML/Qwen-Image-Layered

1import os2import uuid3import numpy as np4import random5import tempfile6import spaces7import zipfile 8from PIL import Image9from diffusers import QwenImageLayeredPipeline10import torch11from pptx import Presentation12import gradio as gr13 14 15LOG_DIR = "/tmp/local"16MAX_SEED = np.iinfo(np.int32).max17 18from huggingface_hub import login19login(token=os.environ.get('hf'))20 21dtype = torch.bfloat1622device = "cuda" if torch.cuda.is_available() else "cpu"23pipeline = QwenImageLayeredPipeline.from_pretrained("Qwen/Qwen-Image-Layered", torch_dtype=dtype).to(device)24# pipeline.set_progress_bar_config(disable=None)25 26def ensure_dirname(path: str):27    if path and not os.path.exists(path):28        os.makedirs(path, exist_ok=True)29 30def random_str(length=8):31    return uuid.uuid4().hex[:length]32 33def imagelist_to_pptx(img_files):34    with Image.open(img_files[0]) as img:35        img_width_px, img_height_px = img.size36 37    def px_to_emu(px, dpi=96):38        inch = px / dpi39        emu = inch * 91440040        return int(emu)41 42    prs = Presentation()43    prs.slide_width = px_to_emu(img_width_px)44    prs.slide_height = px_to_emu(img_height_px)45 46    slide = prs.slides.add_slide(prs.slide_layouts[6])47 48    left = top = 049    for img_path in img_files:50        slide.shapes.add_picture(img_path, left, top, width=px_to_emu(img_width_px), height=px_to_emu(img_height_px))51 52    with tempfile.NamedTemporaryFile(suffix=".pptx", delete=False) as tmp:53        prs.save(tmp.name)54        return tmp.name55 56def export_gallery(images):57    # images: list of image file paths58    images = [e[0] for e in images]59    pptx_path = imagelist_to_pptx(images)60    return pptx_path61 62def export_gallery_zip(images):63    # images: list of tuples (file_path, caption)64    images = [e[0] for e in images]65    66    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:67        with zipfile.ZipFile(tmp.name, 'w', zipfile.ZIP_DEFLATED) as zipf:68            for i, img_path in enumerate(images):69                # Get the file extension from original file70                ext = os.path.splitext(img_path)[1] or '.png'71                # Add each image to the zip with a numbered filename72                zipf.write(img_path, f"layer_{i+1}{ext}")73        return tmp.name74 75@spaces.GPU(duration=180)76def infer(input_image,77          seed=777,78          randomize_seed=False,79          prompt=None,80          neg_prompt=" ",81          true_guidance_scale=4.0,82          num_inference_steps=50,83          layer=4,84          cfg_norm=True,85          use_en_prompt=True):86    87    if randomize_seed:88        seed = random.randint(0, MAX_SEED)89        90    if isinstance(input_image, list):91        input_image = input_image[0]92        93    if isinstance(input_image, str):94        pil_image = Image.open(input_image).convert("RGB").convert("RGBA")95    elif isinstance(input_image, Image.Image):96        pil_image = input_image.convert("RGB").convert("RGBA")97    elif isinstance(input_image, np.ndarray):98        pil_image = Image.fromarray(input_image).convert("RGB").convert("RGBA")99    else:100        raise ValueError("Unsupported input_image type: %s" % type(input_image))101    102    inputs = {103        "image": pil_image,104        "generator": torch.Generator(device='cuda').manual_seed(seed),105        "true_cfg_scale": true_guidance_scale,106        "prompt": prompt,107        "negative_prompt": neg_prompt,108        "num_inference_steps": num_inference_steps,109        "num_images_per_prompt": 1,110        "layers": layer,111        "resolution": 640,      # Using different bucket (640, 1024) to determine the resolution. For this version, 640 is recommended112        "cfg_normalize": cfg_norm,  # Whether enable cfg normalization.113        "use_en_prompt": use_en_prompt, 114    }115    print(inputs)116    with torch.inference_mode():117        output = pipeline(**inputs)118        output_images = output.images[0]119    120    output = []121    temp_files = []122    for i, image in enumerate(output_images):123        output.append(image)124        # Save to temp file for export125        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)126        image.save(tmp.name)127        temp_files.append(tmp.name)128    129    # Generate PPTX130    pptx_path = imagelist_to_pptx(temp_files)131    132    # Generate ZIP133    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:134        with zipfile.ZipFile(tmp.name, 'w', zipfile.ZIP_DEFLATED) as zipf:135            for i, img_path in enumerate(temp_files):136                zipf.write(img_path, f"layer_{i+1}.png")137        zip_path = tmp.name138    139    return output, pptx_path, zip_path140 141ensure_dirname(LOG_DIR)142examples = [143            "assets/test_images/1.png",144            "assets/test_images/2.png",145            "assets/test_images/3.png",146            "assets/test_images/4.png",147            "assets/test_images/5.png",148            "assets/test_images/6.png",149            "assets/test_images/7.png",150            "assets/test_images/8.png",151            "assets/test_images/9.png",152            "assets/test_images/10.png",153            "assets/test_images/11.png",154            "assets/test_images/12.png",155            "assets/test_images/13.png",156            ]157 158 159with gr.Blocks() as demo:160    with gr.Column(elem_id="col-container"):161        gr.HTML('<img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/layered/qwen-image-layered-logo.png" alt="Qwen-Image-Layered Logo" width="600" style="display: block; margin: 0 auto;">')162        gr.Markdown("""163                    The text prompt is intended to describe the overall content of the input image—including elements that may be partially occluded (e.g., you may specify the text hidden behind a foreground object). It is not designed to control the semantic content of individual layers explicitly.164                    """)165        with gr.Row():166            with gr.Column(scale=1):167                input_image = gr.Image(label="Input Image", image_mode="RGBA")168                169                170                with gr.Accordion("Advanced Settings", open=False):171                    prompt = gr.Textbox(172                        label="Prompt (Optional)",173                        placeholder="Please enter the prompt to descibe the image. （Optional）",174                        value="",175                        lines=2,176                    )177                    neg_prompt = gr.Textbox(178                        label="Negative Prompt (Optional)",179                        placeholder="Please enter the negative prompt",180                        value=" ",181                        lines=2,182                    )183                    184                    seed = gr.Slider(185                        label="Seed",186                        minimum=0,187                        maximum=MAX_SEED,188                        step=1,189                        value=0,190                    )191                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)192                    193                    true_guidance_scale = gr.Slider(194                        label="True guidance scale",195                        minimum=1.0,196                        maximum=10.0,197                        step=0.1,198                        value=4.0199                    )200 201                    num_inference_steps = gr.Slider(202                        label="Number of inference steps",203                        minimum=1,204                        maximum=50,205                        step=1,206                        value=50,207                    )208 209                    layer = gr.Slider(210                        label="Layers",211                        minimum=2,212                        maximum=10,213                        step=1,214                        value=4,215                    )216 217                    cfg_norm = gr.Checkbox(label="Whether enable CFG normalization", value=True)218                    use_en_prompt = gr.Checkbox(label="Automatic caption language if no prompt provided, True for EN, False for ZH", value=True)219                220                run_button = gr.Button("Decompose!", variant="primary")221 222            with gr.Column(scale=2):223                gallery = gr.Gallery(label="Layers", columns=4, rows=1, format="png")224                with gr.Row():225                    export_file = gr.File(label="Download PPTX")226                    export_zip_file = gr.File(label="Download ZIP")227 228    gr.Examples(examples=examples,229                inputs=[input_image], 230                outputs=[gallery, export_file, export_zip_file],231                fn=infer, 232                examples_per_page=14,233                cache_examples=False,234                run_on_click=True235    )236 237    run_button.click(238        fn=infer,239        inputs=[240            input_image,241            seed,242            randomize_seed,243            prompt,244            neg_prompt,245            true_guidance_scale,246            num_inference_steps,247            layer,248            cfg_norm,249            use_en_prompt,250        ], 251        outputs=[gallery, export_file, export_zip_file],252    )253 254if __name__ == "__main__":255    demo.launch()