Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import numpy as np | |
| import random | |
| import torch | |
| import spaces | |
| import os | |
| from pathlib import Path | |
| from PIL import Image | |
| from diffusers import FlowMatchEulerDiscreteScheduler | |
| from optimization import optimize_pipeline_ | |
| from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline | |
| from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel | |
| from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3 | |
| import math | |
| # --- Model Loading --- | |
| dtype = torch.bfloat16 | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| scheduler_config = { | |
| "base_image_seq_len": 256, | |
| "base_shift": math.log(3), | |
| "invert_sigmas": False, | |
| "max_image_seq_len": 8192, | |
| "max_shift": math.log(3), | |
| "num_train_timesteps": 1000, | |
| "shift": 1.0, | |
| "shift_terminal": None, | |
| "stochastic_sampling": False, | |
| "time_shift_type": "exponential", | |
| "use_beta_sigmas": False, | |
| "use_dynamic_shifting": True, | |
| "use_exponential_sigmas": False, | |
| "use_karras_sigmas": False, | |
| } | |
| scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config) | |
| pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", scheduler=scheduler, torch_dtype=dtype) | |
| # Load the texture LoRA | |
| pipe.load_lora_weights("2vXpSwA7/iroiro-lora", | |
| weight_name="qwen_lora/qie2509_lora_katame_transferring_01.safetensors", adapter_name="texture") | |
| pipe.load_lora_weights("lightx2v/Qwen-Image-Lightning", | |
| weight_name="Qwen-Image-Lightning-4steps-V2.0-bf16.safetensors", adapter_name="lightning") | |
| pipe.set_adapters(["texture", "lightning"], adapter_weights=[1.2, 1.]) | |
| pipe.fuse_lora(adapter_names=["texture", "lightning"], lora_scale=1) | |
| pipe.unload_lora_weights() | |
| pipe.transformer.__class__ = QwenImageTransformer2DModel | |
| pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3()) | |
| pipe.to(device) | |
| optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt") | |
| MAX_SEED = np.iinfo(np.int32).max | |
| # --- Load sample images --- | |
| def get_sample_images(folder): | |
| """Get all image files from a folder.""" | |
| folder_path = Path(folder) | |
| if not folder_path.exists(): | |
| return [] | |
| image_extensions = {'.png', '.jpg', '.jpeg', '.webp', '.bmp'} | |
| images = [] | |
| for file in sorted(folder_path.iterdir()): | |
| if file.suffix.lower() in image_extensions: | |
| images.append(str(file)) | |
| return images | |
| slotA_images = get_sample_images("samples/slotA") | |
| slotB_images = get_sample_images("samples/slotB") | |
| def calculate_dimensions(image): | |
| """Calculate output dimensions based on content image, keeping largest side at 1024.""" | |
| if image is None: | |
| return 1024, 1024 | |
| original_width, original_height = image.size | |
| if original_width > original_height: | |
| new_width = 1024 | |
| aspect_ratio = original_height / original_width | |
| new_height = int(new_width * aspect_ratio) | |
| else: | |
| new_height = 1024 | |
| aspect_ratio = original_width / original_height | |
| new_width = int(new_height * aspect_ratio) | |
| # Ensure dimensions are multiples of 8 | |
| new_width = (new_width // 8) * 8 | |
| new_height = (new_height // 8) * 8 | |
| return new_width, new_height | |
| def apply_texture( | |
| content_image, | |
| texture_image, | |
| prompt, | |
| seed=42, | |
| randomize_seed=False, | |
| true_guidance_scale=False, | |
| num_inference_steps=4, | |
| progress=gr.Progress(track_tqdm=True) | |
| ): | |
| if content_image is None: | |
| raise gr.Error("Please upload a content image.") | |
| if texture_image is None: | |
| raise gr.Error("Please upload a texture image.") | |
| if not prompt or not prompt.strip(): | |
| prompt = "change image1 character texture to image2 texture" | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| generator = torch.Generator(device=device).manual_seed(seed) | |
| # Calculate dimensions based on content image | |
| width, height = calculate_dimensions(content_image) | |
| # Prepare images | |
| content_pil = content_image.convert("RGB") if isinstance(content_image, Image.Image) else Image.open(content_image.name).convert("RGB") | |
| texture_pil = texture_image.convert("RGB") if isinstance(texture_image, Image.Image) else Image.open(texture_image.name).convert("RGB") | |
| pil_images = [content_pil, texture_pil] | |
| result = pipe( | |
| image=pil_images, | |
| prompt=prompt, | |
| height=height, | |
| width=width, | |
| num_inference_steps=num_inference_steps, | |
| generator=generator, | |
| true_cfg_scale=true_guidance_scale, | |
| num_images_per_prompt=1, | |
| ).images[0] | |
| return result, seed | |
| # --- UI --- | |
| css = ''' | |
| #col-container, #examples { | |
| max-width: 1400px; | |
| margin: 0 auto; | |
| padding: 20px; | |
| } | |
| .dark .progress-text{ | |
| color: white !important; | |
| } | |
| /* Card style for image containers */ | |
| .image-card { | |
| background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); | |
| border-radius: 16px; | |
| padding: 4px; | |
| box-shadow: 0 8px 16px rgba(0,0,0,0.1); | |
| } | |
| /* Input section styling */ | |
| .input-section { | |
| background: rgba(255,255,255,0.05); | |
| border-radius: 12px; | |
| padding: 20px; | |
| margin-bottom: 15px; | |
| } | |
| /* Button styling */ | |
| .generate-btn { | |
| background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important; | |
| border: none !important; | |
| font-size: 18px !important; | |
| font-weight: 600 !important; | |
| padding: 12px 24px !important; | |
| border-radius: 8px !important; | |
| box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4) !important; | |
| transition: all 0.3s ease !important; | |
| } | |
| .generate-btn:hover { | |
| transform: translateY(-2px); | |
| box-shadow: 0 6px 16px rgba(102, 126, 234, 0.6) !important; | |
| } | |
| /* Output section */ | |
| .output-section { | |
| background: rgba(255,255,255,0.03); | |
| border-radius: 12px; | |
| padding: 20px; | |
| min-height: 600px; | |
| } | |
| /* Accordion styling */ | |
| .accordion { | |
| border-radius: 8px; | |
| margin-top: 10px; | |
| } | |
| /* Image upload area */ | |
| .image-upload { | |
| border: 2px dashed rgba(102, 126, 234, 0.3); | |
| border-radius: 12px; | |
| transition: all 0.3s ease; | |
| } | |
| .image-upload:hover { | |
| border-color: rgba(102, 126, 234, 0.6); | |
| background: rgba(102, 126, 234, 0.05); | |
| } | |
| ''' | |
| with gr.Blocks(css=css) as demo: | |
| with gr.Column(elem_id="col-container"): | |
| # Header | |
| gr.Markdown("# 🎨 Qwen Image Edit - Katame Transfer") | |
| gr.Markdown(""" | |
| Transform your images with AI-powered texture transfer using **Qwen Image Edit 2509** | |
| Powered by [2vXpSwA7/iroiro-lora](https://huggingface.co/2vXpSwA7/iroiro-lora) • [Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning) ⚡ | |
| """) | |
| gr.Markdown("---") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### 📥 Input Images") | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("**🖼️ Content Image**") | |
| content_image = gr.Image(label="", type="pil", elem_classes="image-upload") | |
| with gr.Accordion("📁 Sample Images", open=False): | |
| slotA_gallery = gr.Gallery( | |
| value=slotA_images, | |
| label="", | |
| columns=3, | |
| height="auto", | |
| allow_preview=True, | |
| show_label=False | |
| ) | |
| with gr.Column(): | |
| gr.Markdown("**🎨 Texture Image**") | |
| texture_image = gr.Image(label="", type="pil", elem_classes="image-upload") | |
| with gr.Accordion("📁 Sample Textures", open=False): | |
| slotB_gallery = gr.Gallery( | |
| value=slotB_images, | |
| label="", | |
| columns=3, | |
| height="auto", | |
| allow_preview=True, | |
| show_label=False | |
| ) | |
| gr.Markdown("### ✍️ Description") | |
| prompt = gr.Textbox( | |
| label="", | |
| info="", | |
| placeholder="", | |
| lines=2 | |
| ) | |
| button = gr.Button("✨ Generate Image", variant="primary", elem_classes="generate-btn") | |
| with gr.Accordion("⚙️ Advanced Settings", open=False): | |
| seed = gr.Slider(label="🎲 Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) | |
| randomize_seed = gr.Checkbox(label="🔀 Randomize Seed", value=True) | |
| true_guidance_scale = gr.Slider( | |
| label="🎯 Guidance Scale", | |
| minimum=1.0, | |
| maximum=10.0, | |
| step=0.1, | |
| value=1.0, | |
| info="Higher values = stronger adherence to prompt" | |
| ) | |
| num_inference_steps = gr.Slider( | |
| label="⚡ Inference Steps", | |
| minimum=1, | |
| maximum=40, | |
| step=1, | |
| value=4, | |
| info="More steps = higher quality (but slower)" | |
| ) | |
| with gr.Column(scale=1): | |
| gr.Markdown("### 🎭 Generated Result") | |
| output = gr.Image(label="", interactive=False, elem_classes="output-section") | |
| with gr.Row(): | |
| seed_display = gr.Number(label="🌱 Used Seed", interactive=False, visible=True) | |
| # Event handlers | |
| def select_slotA_image(evt: gr.SelectData): | |
| return slotA_images[evt.index] | |
| def select_slotB_image(evt: gr.SelectData): | |
| return slotB_images[evt.index] | |
| slotA_gallery.select(fn=select_slotA_image, outputs=content_image) | |
| slotB_gallery.select(fn=select_slotB_image, outputs=texture_image) | |
| button.click( | |
| fn=apply_texture, | |
| inputs=[ | |
| content_image, | |
| texture_image, | |
| prompt, | |
| seed, | |
| randomize_seed, | |
| true_guidance_scale, | |
| num_inference_steps | |
| ], | |
| outputs=[output, seed_display] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |