Spaces:
Running
on
Zero
Running
on
Zero
| 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 | |
| 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') |