import spaces import os import json import time import torch from PIL import Image from tqdm import tqdm import gradio as gr from safetensors.torch import save_file from src.pipeline import FluxPipeline from src.transformer_flux import FluxTransformer2DModel from src.lora_helper import set_single_lora, set_multi_lora, unset_lora # Initialize the image processor base_path = "black-forest-labs/FLUX.1-dev" lora_base_path = "./models" style_lora_base_path = "Shakker-Labs" pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16) transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat16) pipe.transformer = transformer pipe.to("cuda") def clear_cache(transformer): for name, attn_processor in transformer.attn_processors.items(): attn_processor.bank_kv.clear() # Define the Gradio interface @spaces.GPU() def single_condition_generate_image(prompt, subject_img, spatial_img, height, width, seed, control_type, style_lora=None): # Set the control type if control_type == "subject": lora_path = os.path.join(lora_base_path, "subject.safetensors") elif control_type == "pose": lora_path = os.path.join(lora_base_path, "pose.safetensors") elif control_type == "inpainting": lora_path = os.path.join(lora_base_path, "inpainting.safetensors") set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512) # Set the style LoRA if style_lora=="None": pass else: if style_lora == "Simple_Sketch": pipe.unload_lora_weights() style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Children-Simple-Sketch") pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-children-simple-sketch.safetensors") if style_lora == "Text_Poster": pipe.unload_lora_weights() style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Text-Poster") pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-Text-Poster.safetensors") if style_lora == "Vector_Style": pipe.unload_lora_weights() style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Vector-Journey") pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-Vector-Journey.safetensors") # Process the image subject_imgs = [subject_img] if subject_img else [] spatial_imgs = [spatial_img] if spatial_img else [] image = pipe( prompt, height=int(height), width=int(width), guidance_scale=3.5, num_inference_steps=25, max_sequence_length=512, generator=torch.Generator("cpu").manual_seed(seed), subject_images=subject_imgs, spatial_images=spatial_imgs, cond_size=512, ).images[0] clear_cache(pipe.transformer) return image # Define the Gradio interface @spaces.GPU() def multi_condition_generate_image(prompt, subject_img, spatial_img, height, width, seed): subject_path = os.path.join(lora_base_path, "subject.safetensors") inpainting_path = os.path.join(lora_base_path, "inpainting.safetensors") set_multi_lora(pipe.transformer, [subject_path, inpainting_path], lora_weights=[[1],[1]],cond_size=512) # Process the image subject_imgs = [subject_img] if subject_img else [] spatial_imgs = [spatial_img] if spatial_img else [] image = pipe( prompt, height=int(height), width=int(width), guidance_scale=3.5, num_inference_steps=25, max_sequence_length=512, generator=torch.Generator("cpu").manual_seed(seed), subject_images=subject_imgs, spatial_images=spatial_imgs, cond_size=512, ).images[0] clear_cache(pipe.transformer) return image # Define the Gradio interface components control_types = ["pose", "subject", "inpainting"] style_loras = ["None", "Simple_Sketch", "Text_Poster", "Vector_Style"] # Example data single_examples = [ ["A SKS in the library", Image.open("./test_imgs/subject1.png"), None, 768, 768, 5, "subject", "None"], ["sketched style,A joyful girl with balloons floats above a city wearing a hat and striped pants", None, Image.open("./test_imgs/spatial0.png"), 768, 512, 42, "pose", "Simple_Sketch"], ["In a picturesque village, a narrow cobblestone street with rustic stone buildings, colorful blinds, and lush green spaces, a cartoon man drawn with simple lines and solid colors stands in the foreground, wearing a red shirt, beige work pants, and brown shoes, carrying a strap on his shoulder. The scene features warm and enticing colors, a pleasant fusion of nature and architecture, and the camera's perspective on the street clearly shows the charming and quaint environment., Integrating elements of reality and cartoon.", None, Image.open("./test_imgs/spatial1.png"), 768, 768, 1, "pose", "Vector_Style"], ] multi_examples = [ ["A SKS on the car", Image.open("./test_imgs/subject2.png"), Image.open("./test_imgs/spatial2.png"), 768, 768, 7], ] # Create the Gradio Blocks interface with gr.Blocks() as demo: gr.Markdown("# Image Generation with EasyControl") gr.Markdown("Generate images using EasyControl with different control types and style LoRAs.（Due to hardware constraints, only low-resolution images can be generated. For high-resolution (1024+), please set up your own environment.）") with gr.Tab("Single Condition Generation"): with gr.Row(): with gr.Column(): gr.Markdown(""" **Prompt** (When using LoRA, please try the recommended prompts available at the following links: [FLUX.1-dev-LoRA-Text-Poster](https://huggingface.co/Shakker-Labs/FLUX.1-dev-LoRA-Text-Poster), [FLUX.1-dev-LoRA-Children-Simple-Sketch](https://huggingface.co/Shakker-Labs/FLUX.1-dev-LoRA-Children-Simple-Sketch), [FLUX.1-dev-LoRA-Vector-Journey](https://huggingface.co/Shakker-Labs/FLUX.1-dev-LoRA-Vector-Journey)) """) prompt = gr.Textbox(label="Prompt") subject_img = gr.Image(label="Subject Image", type="pil") # 上传图像文件 spatial_img = gr.Image(label="Spatial Image", type="pil") # 上传图像文件 height = gr.Slider(minimum=256, maximum=1024, step=64, label="Height", value=768) width = gr.Slider(minimum=256, maximum=1024, step=64, label="Width", value=768) seed = gr.Number(label="Seed", value=42) control_type = gr.Dropdown(choices=control_types, label="Control Type") style_lora = gr.Dropdown(choices=style_loras, label="Style LoRA") single_generate_btn = gr.Button("Generate Image") with gr.Column(): single_output_image = gr.Image(label="Generated Image") # Add examples for Single Condition Generation gr.Examples( examples=single_examples, inputs=[prompt, subject_img, spatial_img, height, width, seed, control_type, style_lora], outputs=single_output_image, fn=single_condition_generate_image, cache_examples=False, # 缓存示例结果以加快加载速度 label="Single Condition Examples" ) with gr.Tab("Multi-Condition Generation"): with gr.Row(): with gr.Column(): multi_prompt = gr.Textbox(label="Prompt") multi_subject_img = gr.Image(label="Subject Image", type="pil") # 上传图像文件 multi_spatial_img = gr.Image(label="Spatial Image", type="pil") # 上传图像文件 multi_height = gr.Slider(minimum=256, maximum=1024, step=64, label="Height", value=768) multi_width = gr.Slider(minimum=256, maximum=1024, step=64, label="Width", value=768) multi_seed = gr.Number(label="Seed", value=42) multi_generate_btn = gr.Button("Generate Image") with gr.Column(): multi_output_image = gr.Image(label="Generated Image") # Add examples for Multi-Condition Generation gr.Examples( examples=multi_examples, inputs=[multi_prompt, multi_subject_img, multi_spatial_img, multi_height, multi_width, multi_seed], outputs=multi_output_image, fn=multi_condition_generate_image, cache_examples=False, # 缓存示例结果以加快加载速度 label="Multi-Condition Examples" ) # Link the buttons to the functions single_generate_btn.click( single_condition_generate_image, inputs=[prompt, subject_img, spatial_img, height, width, seed, control_type, style_lora], outputs=single_output_image ) multi_generate_btn.click( multi_condition_generate_image, inputs=[multi_prompt, multi_subject_img, multi_spatial_img, multi_height, multi_width, multi_seed], outputs=multi_output_image ) # Launch the Gradio app demo.queue().launch()