app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces

from PIL import Image
from diffsynth.pipelines.qwen_image import QwenImagePipeline, ModelConfig

import os

from huggingface_hub import  hf_hub_download



pipe = QwenImagePipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
        ModelConfig(model_id="Qwen/Qwen-Image-Edit-2509", 
        download_source='huggingface',
        origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
        ModelConfig(model_id="Qwen/Qwen-Image-Edit-2509", 
        download_source='huggingface',origin_file_pattern="text_encoder/model*.safetensors"),
        ModelConfig(model_id="Qwen/Qwen-Image-Edit-2509", 
        download_source='huggingface',origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
    ],
    tokenizer_config=None,
    processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Edit-2509", 
    download_source='huggingface',origin_file_pattern="processor/"),
)






speedup = hf_hub_download(repo_id="Tele-AI/TeleStyle", filename="weights/diffsynth_Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors")
telestyle= hf_hub_download(repo_id="Tele-AI/TeleStyle", filename="weights/diffsynth_Qwen-Image-Edit-2509-telestyle.safetensors")


pipe.load_lora(pipe.dit, telestyle)
pipe.load_lora(pipe.dit,speedup)



dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"




MAX_SEED = np.iinfo(np.int32).max


@spaces.GPU
def infer(
    content_ref,
    style_ref,
    prompt,
    seed=123,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    minedge=1024,
    progress=gr.Progress(track_tqdm=True),
    
):
    
    

    

    content_ref=Image.fromarray(content_ref)
    style_ref=Image.fromarray(style_ref)
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    
    
    
    
    
    
    w,h=content_ref.size



    minedge=minedge-minedge%16
    if w>h:
        r=w/h
        h=minedge
        w=int(h*r)-int(h*r)%16
        
    else:
        r=h/w
        w=minedge
        h=int(w*r)-int(w*r)%16


    
    print(f"Calling pipeline with prompt: '{prompt}'")
    
    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {w}x{h}")
    
    images = [
        content_ref.resize((w, h)),
        style_ref.resize((minedge, minedge)) ,
    ]
    

    # Generate the image
    
    image = pipe(prompt, edit_image=images, seed=seed, num_inference_steps=num_inference_steps, height=h, width=w,edit_image_auto_resize=False,cfg_scale=true_guidance_scale)#ligtning


    return image, seed

# --- Examples and UI Layout ---
examples = []



_HEADER_ = '''
<div style="text-align: center; max-width: 650px; margin: 0 auto;">
    <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">TeleStyle</h1>
    
</div>


<p style="font-size: 1rem; margin-bottom: 1.5rem;">Paper: <a href='https://arxiv.org/abs/2601.20175' target='_blank'>TeleStyle: Content-Preserving Style Transfer in Images and Videos</a> | Codes: <a href='https://github.com/Tele-AI/TeleStyle/' target='_blank'>GitHub</a></p>
<p style="font-size: 1rem; margin-bottom: 1.5rem;">If you encounter an Error with this demo, the most possible reason is ZeroGPU out-of-memory and the solution is to decrease the Min Edge of the generated image from 1024 to a lower value. This is because ZeroGPU has a memory limit of 70GB, while all the examples are tested with 80GB H100 GPUs.  </p>
'''  

with gr.Blocks() as demo:

    with gr.Column(elem_id="col-container"):
        
        gr.Markdown(_HEADER_)
        gr.Markdown("This is a demo of TeleStyle-Image, enabling Content-Preserving Style Transfer capability to Qwen-Image-Edit-2509.")
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    content_ref = gr.Image(label="content ref", type="numpy", )
                    style_ref = gr.Image(label="style ref", type="numpy", )
                    #print(f"type(content_ref)={type(content_ref)}")
                    
                #input_images = gr.Gallery(label="Input Images", show_label=False, type="pil", interactive=True)
            

            result = gr.Image(label="Result", show_label=True, type="pil")
            #result = gr.Gallery(label="Result", show_label=True, type="pil")
        with gr.Row():
            prompt = gr.Text(
                    label="Prompt",
                    value='Style Transfer the style of Figure 2 to Figure 1, and keep the content and characteristics of Figure 1.',
                    show_label=True,
                    placeholder='Style Transfer the style of Figure 2 to Figure 1, and keep the content and characteristics of Figure 1.',
                    container=True,
            )
            run_button = gr.Button("Edit!", variant="primary")

        with gr.Accordion("Advanced Settings", open=True):
            # Negative prompt UI element is removed here

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=123,
            )

            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)

            with gr.Row():

                true_guidance_scale = gr.Slider(
                    label="CFG should be 1.0",
                    minimum=0,
                    maximum=10.0,
                    step=0.1,
                    value=1.0
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps should be 4",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )
                
                minedge = gr.Slider(
                    label="Min Edge of the generated image",
                    minimum=256,
                    maximum=2048,
                    step=8,
                    value=1024,
                )
        
    
        with gr.Row(), gr.Column():
            gr.Markdown("## Examples")
            gr.Markdown("changing the minedge could lead to different style similarity.")
            default_prompt='Style Transfer the style of Figure 2 to Figure 1, and keep the content and characteristics of Figure 1.'
            gr.Examples(examples=[
                ['./qwenstyleref/pulpfiction_2.jpg','./qwenstyleref/styleref=6_style_ref.png',default_prompt,123,False,1.0,4,832],
                ['./qwenstyleref/styleref=0_content_ref.png','./qwenstyleref/110.png',default_prompt,123,False,1.0,4,832],
                ['./qwenstyleref/romanholiday_1.jpg','./qwenstyleref/s0099____1113_01_query_1_img_000146_1682705733350_08158389675901344.jpg.jpg',default_prompt,123,False,1.0,4,800],
                ['./qwenstyleref/styleref=0_content_ref.png','./qwenstyleref/125.png',default_prompt,123,False,1.0,4,832],
                ['./qwenstyleref/fallenangle.jpg','./qwenstyleref/styleref=s0038.png',default_prompt,123,False,1.0,4,832],
                ['./qwenstyleref/styleref=0_content_ref.png','./qwenstyleref/styleref=s0572.png',default_prompt,123,False,1.0,4,832],
                ['./qwenstyleref/startrooper1.jpg','./qwenstyleref/david-face-760x985.jpg','Style Transfer Figure  1 into marble material.',123,False,1.0,4,1024],
                ['./qwenstyleref/startrooper1.jpg','./qwenstyleref/125.png',default_prompt, 123,False,1.0,4,1024],
                ['./qwenstyleref/possession.png','./qwenstyleref/s0026____0907_01_query_0_img_000194_1682674358294_041656249089406583.jpeg.jpg',default_prompt,123,False,1.0,4,832],
                ['./qwenstyleref/styleref=0_content_ref.png','./qwenstyleref/Jotarokujo.webp',default_prompt,123,False,1.0,4,832],
                ['./qwenstyleref/wallstreet1.jpg','./qwenstyleref/034.png',default_prompt,123,False,1.0,4,1024],
                ['./qwenstyleref/bird.jpeg','./qwenstyleref/styleref=s0539.png',default_prompt,123,False,1.0,4,832],
                

                
                ],
                inputs=[content_ref,
                    style_ref,
                    prompt,
                    seed,
                    randomize_seed,
                    true_guidance_scale,
                    num_inference_steps,
                    minedge,],
                #inputs=[content_ref,style_ref, prompt,], 
                outputs=[result, seed], 
                fn=infer, 
                cache_examples=False
                )        
                
                
                
                

        # gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)

    gr.on(
        triggers=[run_button.click],
        fn=infer,
        inputs=[
            content_ref,
            style_ref,
            prompt,
            seed,
            randomize_seed,
            true_guidance_scale,
            num_inference_steps,
            minedge,
            
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch(server_name='0.0.0.0')