import gradio as gr import numpy as np import torch, random, json, spaces, time from ulid import ULID from diffsynth.pipelines.qwen_image import ( QwenImagePipeline as QwenImagePipelineDs, ModelConfig, QwenImageUnit_Image2LoRAEncode, QwenImageUnit_Image2LoRADecode ) from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler from transformers import Qwen2_5_VLForConditionalGeneration from diffusers import QwenImagePipeline, QwenImageTransformer2DModel from safetensors.torch import save_file import torch, math from PIL import Image DTYPE = torch.bfloat16 MAX_SEED = np.iinfo(np.int32).max from torchao.quantization import ( quantize_, Float8WeightOnlyConfig, Float8DynamicActivationFloat8WeightConfig, FqnToConfig ) pipe_imagen = QwenImagePipeline.from_pretrained( "Qwen/Qwen-Image", torch_dtype=torch.bfloat16 ).to("cuda") print(f'Merging Loras') pipe_imagen.load_lora_weights( "lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-4steps-V2.0-bf16.safetensors", adapter_name="lightning_steps" ) pipe_imagen.set_adapters(["lightning_steps"], adapter_weights=[1.0]) pipe_imagen.fuse_lora(adapter_names=["lightning_steps"], lora_scale=1.25) pipe_imagen.unload_lora_weights() fp8quant = Float8WeightOnlyConfig() transformer_qconfig_dict = { "re:.*bias.*": None, "re:.*proj.*weight": fp8quant, "re:.*to_add_out.*weight": fp8quant, "re:.*to_k.*weight": fp8quant, "re:.*to_out.*weight": fp8quant, "re:.*to_q.*weight": fp8quant, "re:.*to_v.*weight": fp8quant, "re:.*txt_mod\\.1.*weight": fp8quant, "re:.*img_mod\\.1.*weight": fp8quant, # loras "re:.*alpha": None, "re:.*lora.*": None, "_default": fp8quant, } text_encoder_qconfig_dict = { "re:.*bias.*": None, # "re:.*proj.*weight": fp8quant, "re:.*to_add_out.*weight": fp8quant, "re:.*to_k.*weight": fp8quant, "re:.*to_out.*weight": fp8quant, "re:.*to_q.*weight": fp8quant, "re:.*to_v.*weight": fp8quant, "_default": fp8quant, } print('Quantize transformer') quantize_(pipe_imagen.transformer, FqnToConfig(transformer_qconfig_dict), filter_fn=None) print('Quantize text_encoder') quantize_(pipe_imagen.text_encoder, FqnToConfig(text_encoder_qconfig_dict), filter_fn=None) # quant_model = "AiSudo/Qwen-Image-fp8-4steps" # print('Loading transformer...') # transformer=QwenImageTransformer2DModel.from_pretrained( # quant_model, # subfolder='transformer', # torch_dtype=torch.bfloat16, # use_safetensors=False, # device_map="cuda" # ) # print('Loading text_encoder...') # text_encoder=Qwen2_5_VLForConditionalGeneration.from_pretrained( # quant_model, # subfolder='text_encoder', # torch_dtype=torch.bfloat16, # use_safetensors=False, # device_map="cuda" # ) # scheduler = FlowMatchEulerDiscreteScheduler.from_config( # quant_model, # subfolder="scheduler" # ) # pipe_imagen = QwenImagePipeline.from_pretrained( # "Qwen/Qwen-Image", # transformer=transformer, # text_encoder=text_encoder, # scheduler=scheduler, # torch_dtype=torch.bfloat16, # low_cpu_mem_usage=False, # ).to("cuda") # print('Pipe imagen loaded!') # vram_config_disk_offload = { # "offload_dtype": "disk", # "offload_device": "disk", # "onload_dtype": "disk", # "onload_device": "disk", # "preparing_dtype": torch.bfloat16, # "preparing_device": "cuda", # "computation_dtype": torch.bfloat16, # "computation_device": "cuda", # } vram_config_disk_offload = { "offload_dtype": torch.bfloat16, "offload_device": "cuda", "onload_dtype": torch.bfloat16, "onload_device": "cuda", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, "computation_device": "cuda", } # Load models pipe_lora = QwenImagePipelineDs.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ ModelConfig( download_source="huggingface", model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="SigLIP2-G384/model.safetensors", **vram_config_disk_offload ), ModelConfig( download_source="huggingface", model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="DINOv3-7B/model.safetensors", **vram_config_disk_offload ), ModelConfig( download_source="huggingface", model_id="DiffSynth-Studio/Qwen-Image-i2L", origin_file_pattern="Qwen-Image-i2L-Style.safetensors", **vram_config_disk_offload ), ], processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Edit", origin_file_pattern="processor/"), vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) vram_config = { "offload_dtype": "disk", "offload_device": "disk", "onload_dtype": torch.bfloat16, "onload_device": "cuda", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, "computation_device": "cuda", } # pipe_imagen = QwenImagePipelineDs.from_pretrained( # torch_dtype=torch.bfloat16, # device="cuda", # model_configs=[ # ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors", **vram_config), # ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors", **vram_config), # ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config), # ], # tokenizer_config=ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"), # vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, # ) @spaces.GPU def generate_lora( input_images, progress=gr.Progress(track_tqdm=True), ): ulid = str(ULID()).lower()[:12] print(f"ulid: {ulid}") if not input_images: raise gr.Error("images are empty.") input_images = [Image.open(filepath).convert("RGB") for filepath, _ in input_images] # Model inference with torch.no_grad(): embs = QwenImageUnit_Image2LoRAEncode().process(pipe_lora, image2lora_images=input_images) lora = QwenImageUnit_Image2LoRADecode().process(pipe_lora, **embs)["lora"] lora_name = f"{ulid}.safetensors" lora_path = f"loras/{lora_name}" save_file(lora, lora_path) return lora_name, gr.update(interactive=True, value=lora_path), gr.update(interactive=True) @spaces.GPU def generate_image( lora_name, prompt, negative_prompt="blurry ugly bad", width=1024, height=1024, seed=42, randomize_seed=True, guidance_scale=1.0, lora_strength = 1.25, num_inference_steps=8, progress=gr.Progress(track_tqdm=True), ): # lora_path = f"loras/{lora_name}" # pipe_imagen.clear_lora() # pipe_imagen.load_lora(pipe_imagen.dit, lora_path) pipe_imagen.unload_lora_weights() pipe_imagen.load_lora_weights( "loras", weight_name=lora_name, adapter_name="generated_lora" ) pipe_imagen.set_adapters("generated_lora", adapter_weights=lora_strength) if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator().manual_seed(seed) output_image = pipe_imagen( prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, width=width, height=height, true_cfg_scale=guidance_scale, generator=generator, ).images[0] return output_image, seed def read_file(path: str) -> str: with open(path, 'r', encoding='utf-8') as f: content = f.read() return content css = """ #col-container { margin: 0 auto; max-width: 960px; } h3{ text-align: center; display:block; } """ with open('examples/0_examples.json', 'r') as file: examples = json.load(file) print(examples) with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): with gr.Column(): gr.HTML(read_file("static/header.html")) with gr.Row(): with gr.Column(): input_images = gr.Gallery( label="Input images", file_types=["image"], show_label=False, elem_id="gallery", columns=2, object_fit="cover", height=300) lora_button = gr.Button("Generate LoRA", variant="primary") with gr.Column(): lora_name = gr.Textbox(label="Generated LoRA path",lines=2, interactive=False) lora_download = gr.DownloadButton(label=f"Download LoRA", interactive=False) with gr.Column(elem_id='imagen-container') as imagen_container: gr.Markdown("### After your LoRA is ready, you can try generate image here.") with gr.Row(): with gr.Column(): prompt = gr.Textbox( label="Prompt", show_label=False, lines=2, placeholder="Enter your prompt", value="a man in a fishing boat.", container=False, ) lora_strength = gr.Slider( label="Lora Strength", minimum=0.0, maximum=3.0, step=0.01, value=0.85, ) imagen_button = gr.Button("Generate Image", variant="primary", interactive=False) with gr.Accordion("Advanced Settings", open=False): negative_prompt = gr.Textbox( label="Negative prompt", lines=2, container=False, placeholder="Enter your negative prompt", value="blurry ugly bad" ) num_inference_steps = gr.Slider( label="Steps", minimum=2, maximum=12, step=1, value=4, ) with gr.Row(): width = gr.Slider( label="Width", minimum=512, maximum=2048, step=32, value=1024, ) height = gr.Slider( label="Height", minimum=512, maximum=2048, step=32, value=1024, ) with gr.Row(): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) guidance_scale = gr.Slider( label="Guidance scale", minimum=0.0, maximum=4.0, step=0.1, value=0.0, ) with gr.Column(): output_image = gr.Image(label="Generated image", show_label=False) gr.Examples(examples=examples, inputs=[input_images]) gr.Markdown(read_file("static/footer.md")) lora_button.click( fn=generate_lora, inputs=[ input_images ], outputs=[lora_name, lora_download, imagen_button], ) imagen_button.click( fn=generate_image, inputs=[ lora_name, prompt, negative_prompt, width, height, seed, randomize_seed, guidance_scale, lora_strength, num_inference_steps, ], outputs=[output_image, seed], ) if __name__ == "__main__": demo.launch(mcp_server=True, css=css)