app.py

import spaces
import gradio as gr
import torch
from PIL import Image
from diffusers import DiffusionPipeline, QwenImageEditPipeline, FlowMatchEulerDiscreteScheduler
import random
import uuid
import numpy as np
import time
import zipfile
import os
import requests
from urllib.parse import urlparse
import tempfile
import shutil
import math

# --- App Description ---
DESCRIPTION = """## Qwen Image Hpc/."""

# --- Helper Functions for Both Tabs ---
MAX_SEED = np.iinfo(np.int32).max

def save_image(img):
    """Saves a PIL image to a temporary file with a unique name."""
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    """Returns a random seed if randomize_seed is True, otherwise returns the original seed."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# --- Qwen-Image-Gen Model ---
pipe_qwen_gen = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image",
    torch_dtype=dtype
).to(device)

# --- Qwen-Image-Edit Model with Lightning LoRA ---
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_qwen_edit = QwenImageEditPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit",
    scheduler=scheduler,
    torch_dtype=dtype
).to(device)

try:
    pipe_qwen_edit.load_lora_weights(
        "lightx2v/Qwen-Image-Lightning",
        weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors"
    )
    pipe_qwen_edit.fuse_lora()
    print("Successfully loaded Lightning LoRA weights for Qwen-Image-Edit")
except Exception as e:
    print(f"Warning: Could not load Lightning LoRA weights for Qwen-Image-Edit: {e}")
    print("Continuing with the base Qwen-Image-Edit model...")


# --- Qwen-Image-Gen Functions ---
aspect_ratios = {
    "1:1": (1328, 1328),
    "16:9": (1664, 928),
    "9:16": (928, 1664),
    "4:3": (1472, 1140),
    "3:4": (1140, 1472)
}

def load_lora_opt(pipe, lora_input):
    """Loads a LoRA from a local path, Hugging Face repo, or URL."""
    lora_input = lora_input.strip()
    if not lora_input:
        return

    if "/" in lora_input and not lora_input.startswith("http"):
        pipe.load_lora_weights(lora_input, adapter_name="default")
        return

    if lora_input.startswith("http"):
        url = lora_input
        if "huggingface.co" in url and "/blob/" not in url and "/resolve/" not in url:
            repo_id = urlparse(url).path.strip("/")
            pipe.load_lora_weights(repo_id, adapter_name="default")
            return

        if "/blob/" in url:
            url = url.replace("/blob/", "/resolve/")

        tmp_dir = tempfile.mkdtemp()
        local_path = os.path.join(tmp_dir, os.path.basename(urlparse(url).path))

        try:
            print(f"Downloading LoRA from {url}...")
            resp = requests.get(url, stream=True)
            resp.raise_for_status()
            with open(local_path, "wb") as f:
                for chunk in resp.iter_content(chunk_size=8192):
                    f.write(chunk)
            print(f"Saved LoRA to {local_path}")
            pipe.load_lora_weights(local_path, adapter_name="default")
        finally:
            shutil.rmtree(tmp_dir, ignore_errors=True)

@spaces.GPU(duration=120)
def generate_qwen(
    prompt: str,
    negative_prompt: str = "",
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 4.0,
    randomize_seed: bool = False,
    num_inference_steps: int = 50,
    num_images: int = 1,
    zip_images: bool = False,
    lora_input: str = "",
    lora_scale: float = 1.0,
    progress=gr.Progress(track_tqdm=True),
):
    """Main generation function for Qwen-Image-Gen."""
    seed = randomize_seed_fn(seed, randomize_seed)
    generator = torch.Generator(device).manual_seed(seed)

    start_time = time.time()

    current_adapters = pipe_qwen_gen.get_list_adapters()
    for adapter in current_adapters:
        pipe_qwen_gen.delete_adapters(adapter)
    pipe_qwen_gen.disable_lora()

    if lora_input and lora_input.strip() != "":
        load_lora_opt(pipe_qwen_gen, lora_input)
        pipe_qwen_gen.set_adapters(["default"], adapter_weights=[lora_scale])

    images = pipe_qwen_gen(
        prompt=prompt,
        negative_prompt=negative_prompt if negative_prompt else " ",
        height=height,
        width=width,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        num_images_per_prompt=num_images,
        generator=generator,
    ).images

    end_time = time.time()
    duration = end_time - start_time

    image_paths = [save_image(img) for img in images]
    zip_path = None
    if zip_images and len(image_paths) > 0:
        zip_name = str(uuid.uuid4()) + ".zip"
        with zipfile.ZipFile(zip_name, 'w') as zipf:
            for i, img_path in enumerate(image_paths):
                zipf.write(img_path, arcname=f"Img_{i}.png")
        zip_path = zip_name

    current_adapters = pipe_qwen_gen.get_list_adapters()
    for adapter in current_adapters:
        pipe_qwen_gen.delete_adapters(adapter)
    pipe_qwen_gen.disable_lora()

    return image_paths, seed, f"{duration:.2f}", zip_path

@spaces.GPU(duration=120)
def generate(
    prompt: str,
    negative_prompt: str,
    use_negative_prompt: bool,
    seed: int,
    width: int,
    height: int,
    guidance_scale: float,
    randomize_seed: bool,
    num_inference_steps: int,
    num_images: int,
    zip_images: bool,
    lora_input: str,
    lora_scale: float,
    progress=gr.Progress(track_tqdm=True),
):
    """UI wrapper for the Qwen-Image-Gen generation function."""
    final_negative_prompt = negative_prompt if use_negative_prompt else ""
    return generate_qwen(
        prompt=prompt,
        negative_prompt=final_negative_prompt,
        seed=seed,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        randomize_seed=randomize_seed,
        num_inference_steps=num_inference_steps,
        num_images=num_images,
        zip_images=zip_images,
        lora_input=lora_input,
        lora_scale=lora_scale,
        progress=progress,
    )

# --- Qwen-Image-Edit Functions ---
@spaces.GPU(duration=60)
def infer_edit(
    image,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=8,
    progress=gr.Progress(track_tqdm=True),
):
    """Main inference function for Qwen-Image-Edit."""
    if image is None:
        raise gr.Error("Please upload an image to edit.")

    negative_prompt = " "
    seed = randomize_seed_fn(seed, randomize_seed)
    generator = torch.Generator(device=device).manual_seed(seed)

    print(f"Original prompt: '{prompt}'")
    print(f"Negative Prompt: '{negative_prompt}'")
    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")

    try:
        images = pipe_qwen_edit(
            image,
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            generator=generator,
            true_cfg_scale=true_guidance_scale,
            num_images_per_prompt=1
        ).images
        return images[0], seed
    except Exception as e:
        print(f"Error during inference: {e}")
        raise gr.Error(f"An error occurred during image editing: {e}")

# --- Gradio UI ---
css = '''
.gradio-container {
    max-width: 800px !important;
    margin: 0 auto !important;
}
h1 {
    text-align: center;
}
footer {
    visibility: hidden;
}
'''

with gr.Blocks(css=css, theme="bethecloud/storj_theme", delete_cache=(240, 240)) as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Tabs():
        with gr.TabItem("Qwen-Image-Gen"):
            with gr.Column():
                with gr.Row():
                    prompt_gen = gr.Text(
                        label="Prompt",
                        show_label=False,
                        max_lines=1,
                        placeholder="✦︎ Enter your prompt for generation",
                        container=False,
                    )
                    run_button_gen = gr.Button("Generate", scale=0, variant="primary")
                result_gen = gr.Gallery(label="Result", columns=2, show_label=False, preview=True, height="auto")

                with gr.Row():
                    aspect_ratio_gen = gr.Dropdown(
                        label="Aspect Ratio",
                        choices=list(aspect_ratios.keys()),
                        value="1:1",
                    )
                    lora_gen = gr.Textbox(label="Optional LoRA", placeholder="Enter Hugging Face repo ID or URL...")

                with gr.Accordion("Additional Options", open=False):
                    use_negative_prompt_gen = gr.Checkbox(label="Use negative prompt", value=True)
                    negative_prompt_gen = gr.Text(
                        label="Negative prompt",
                        max_lines=1,
                        placeholder="Enter a negative prompt",
                        value="text, watermark, copyright, blurry, low resolution",
                    )
                    seed_gen = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    randomize_seed_gen = gr.Checkbox(label="Randomize seed", value=True)
                    with gr.Row():
                        width_gen = gr.Slider(label="Width", minimum=512, maximum=2048, step=64, value=1328)
                        height_gen = gr.Slider(label="Height", minimum=512, maximum=2048, step=64, value=1328)
                    guidance_scale_gen = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=20.0, step=0.1, value=4.0)
                    num_inference_steps_gen = gr.Slider("Number of inference steps", 1, 100, 50, step=1)
                    num_images_gen = gr.Slider("Number of images", 1, 5, 1, step=1)
                    zip_images_gen = gr.Checkbox(label="Zip generated images", value=False)
                    with gr.Row():
                        lora_scale_gen = gr.Slider(label="LoRA Scale", minimum=0, maximum=2, step=0.01, value=1)

                    gr.Markdown("### Output Information")
                    seed_display_gen = gr.Textbox(label="Seed used", interactive=False)
                    generation_time_gen = gr.Textbox(label="Generation time (seconds)", interactive=False)
                    zip_file_gen = gr.File(label="Download ZIP")
            
            # --- Gen Tab Logic ---
            def set_dimensions(ar):
                w, h = aspect_ratios[ar]
                return gr.update(value=w), gr.update(value=h)

            aspect_ratio_gen.change(fn=set_dimensions, inputs=aspect_ratio_gen, outputs=[width_gen, height_gen])
            use_negative_prompt_gen.change(fn=lambda x: gr.update(visible=x), inputs=use_negative_prompt_gen, outputs=negative_prompt_gen)

            gen_inputs = [
                prompt_gen, negative_prompt_gen, use_negative_prompt_gen, seed_gen, width_gen, height_gen,
                guidance_scale_gen, randomize_seed_gen, num_inference_steps_gen, num_images_gen,
                zip_images_gen, lora_gen, lora_scale_gen
            ]
            gen_outputs = [result_gen, seed_display_gen, generation_time_gen, zip_file_gen]

            gr.on(triggers=[prompt_gen.submit, run_button_gen.click], fn=generate, inputs=gen_inputs, outputs=gen_outputs)

            gen_examples = [
                "A decadent slice of layered chocolate cake on a ceramic plate with a drizzle of chocolate syrup and powdered sugar dusted on top.",
                "A young girl wearing school uniform stands in a classroom, writing on a chalkboard. The text 'Introducing Qwen-Image' appears in neat white chalk.",
                "一幅精致细腻的工笔画，画面中心是一株蓬勃生长的红色牡丹，花朵繁茂。",
                "Realistic still life photography style: A single, fresh apple, resting on a clean, soft-textured surface.",
            ]
            gr.Examples(examples=gen_examples, inputs=prompt_gen, outputs=gen_outputs, fn=generate, cache_examples=False)

        with gr.TabItem("Qwen-Image-Edit"):
            with gr.Column():
                with gr.Row():
                    input_image_edit = gr.Image(label="Input Image", type="pil", height=400)
                    result_edit = gr.Image(label="Result", type="pil", height=400)
                
                with gr.Row():
                    prompt_edit = gr.Text(
                        label="Edit Instruction",
                        show_label=False,
                        placeholder="Describe the edit you want to make",
                        container=False,
                    )
                    run_button_edit = gr.Button("Edit", variant="primary")

                with gr.Accordion("Advanced Settings", open=False):
                    seed_edit = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
                    randomize_seed_edit = gr.Checkbox(label="Randomize seed", value=True)
                    with gr.Row():
                        true_guidance_scale_edit = gr.Slider(
                            label="True guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0
                        )
                        num_inference_steps_edit = gr.Slider(
                            label="Inference steps (Lightning LoRA)", minimum=4, maximum=28, step=1, value=8
                        )
            
            # --- Edit Tab Logic ---
            edit_inputs = [
                input_image_edit, prompt_edit, seed_edit, randomize_seed_edit,
                true_guidance_scale_edit, num_inference_steps_edit
            ]
            edit_outputs = [result_edit, seed_edit]
            
            gr.on(triggers=[prompt_edit.submit, run_button_edit.click], fn=infer_edit, inputs=edit_inputs, outputs=edit_outputs)
            
            edit_examples = [
                ["image-edit/cat.png", "make the cat wear sunglasses"],
                ["image-edit/girl.png", "change her hair to blonde"],
            ]
       
            gr.Examples(examples=edit_examples, inputs=[input_image_edit, prompt_edit], outputs=edit_outputs, fn=infer_edit, cache_examples=True)


if __name__ == "__main__":
    demo.queue(max_size=50).launch(share=False, debug=True)