app.py

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

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("tarn59/apply_texture_qwen_image_edit_2509", 
        weight_name="apply_texture_v2_qwen_image_edit_2509.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., 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

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

@spaces.GPU
def apply_texture(
    content_image,
    texture_image,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    progress=gr.Progress(track_tqdm=True)
):
    if content_image is None:
        raise gr.Error("콘텐츠 이미지를 업로드해주세요.")
    if texture_image is None:
        raise gr.Error("텍스처 이미지를 업로드해주세요.")
    if not prompt or not prompt.strip():
        raise gr.Error("설명을 입력해주세요.")
    
    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: 1200px; margin: 0 auto; }
.dark .progress-text{color: white !important}
'''

with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# 텍스처 적용 — Qwen Image Edit")
        gr.Markdown("""
            [tarn59의 Apply-Texture-Qwen-Image-Edit-2509 LoRA](https://huggingface.co/tarn59/apply_texture_qwen_image_edit_2509)와
            [lightx2v/Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning)을 사용한 4단계 추론 💨
        """)
        
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    content_image = gr.Image(label="콘텐츠", type="pil")
                    texture_image = gr.Image(label="텍스처", type="pil")

                prompt = gr.Textbox(
                    label="설명",
                    info="...에 ... 텍스처 적용",
                    placeholder="건물 벽에 나무 사이딩 텍스처 적용"
                )

                button = gr.Button("✨ 생성", variant="primary")
                
                with gr.Accordion("⚙️ 고급 설정", open=False):
                    seed = gr.Slider(label="시드", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    randomize_seed = gr.Checkbox(label="시드 무작위화", value=True)
                    true_guidance_scale = gr.Slider(
                        label="실제 가이던스 스케일",
                        minimum=1.0,
                        maximum=10.0,
                        step=0.1,
                        value=1.0
                    )
                    num_inference_steps = gr.Slider(
                        label="추론 단계",
                        minimum=4,
                        maximum=40,
                        step=1,
                        value=4
                    )
            
            with gr.Column():
                output = gr.Image(label="출력", interactive=False)
        
        # Event handlers
        button.click(
            fn=apply_texture,
            inputs=[
                content_image,
                texture_image,
                prompt,
                seed,
                randomize_seed,
                true_guidance_scale,
                num_inference_steps
            ],
            outputs=[output, seed]
        )
        
        # Examples
        gr.Examples(
            examples=[
                ["coffee_mug.png", "wood_boxes.png", "머그에 나무 텍스처 적용"],
                ["720park.jpg", "black-and-white.jpg", "건물에 흑백 물결 텍스처 적용"],
                ["red_apple.png", "marble.png", "사과에 흰색-회색 대리석 텍스처 적용"],
                ["sports_car.png", "bamboo.png", "스포츠카에 대나무 텍스처 적용"],
                ["cat_statue.png", "pink_fur.png", "고양이 조각상에 분홍색 털 텍스처 적용"],
            ],
            inputs=[
                content_image,
                texture_image,
                prompt,
            ],
            outputs=[output, seed],
            fn=apply_texture,
            cache_examples="lazy",
            elem_id="examples"
        )

if __name__ == "__main__":
    demo.launch()