app.py

import os
import json
import copy
import math
import time
import random
import logging
import numpy as np
from typing import Any, Dict, List, Optional, Union
import torch
from PIL import Image
import gradio as gr
import spaces
from diffusers import (
    DiffusionPipeline,
    AutoencoderKL,
    AutoencoderTiny,
    AutoPipelineForImage2Image,
    FlowMatchEulerDiscreteScheduler
)
from huggingface_hub import (
    hf_hub_download,
    HfFileSystem,
    ModelCard,
    snapshot_download
)
from diffusers.utils import load_image
import requests
from urllib.parse import urlparse
import tempfile
import shutil
import uuid
import zipfile


# META: CUDA_CHECK / GPU_INFO
device = "cuda" if torch.cuda.is_available() else "cpu"
print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("torch.__version__ =", torch.__version__)
print("torch.version.cuda =", torch.version.cuda)
print("cuda available:", torch.cuda.is_available())
print("cuda device count:", torch.cuda.device_count())
if torch.cuda.is_available():
    print("current device:", torch.cuda.current_device())
    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))

print("Using device:", device)

loras = [
    # Sample Qwen-compatible LoRAs
    {
        "image": "https://huggingface.co/damnthatai/Game_Boy_Camera_Pixel_Style_Qwen/resolve/main/images/20250818090201_Qwen8s_00001_.jpg",
        "title": "Camera Pixel Style",
        "repo": "damnthatai/Game_Boy_Camera_Pixel_Style_Qwen",
        "weights": "g4m3b0yc4m3r4_qwen.safetensors",
        "trigger_word": "g4m3b0yc4m3r4, grayscale, pixel photo"
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Studio-Realism/resolve/main/images/2.png",
        "title": "Studio Realism",
        "repo": "prithivMLmods/Qwen-Image-Studio-Realism",
        "weights": "qwen-studio-realism.safetensors",
        "trigger_word": "Studio Realism"
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Sketch-Smudge/resolve/main/images/1.png",
        "title": "Sketch Smudge",
        "repo": "prithivMLmods/Qwen-Image-Sketch-Smudge",
        "weights": "qwen-sketch-smudge.safetensors",
        "trigger_word": "Sketch Smudge"
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Anime-LoRA/resolve/main/images/1.png",
        "title": "Qwen Anime",
        "repo": "prithivMLmods/Qwen-Image-Anime-LoRA",
        "weights": "qwen-anime.safetensors",
        "trigger_word": "Qwen Anime"
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Fragmented-Portraiture/resolve/main/images/3.png",
        "title": "Fragmented Portraiture",
        "repo": "prithivMLmods/Qwen-Image-Fragmented-Portraiture",
        "weights": "qwen-fragmented-portraiture.safetensors",
        "trigger_word": "Fragmented Portraiture"
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Synthetic-Face/resolve/main/images/2.png",
        "title": "Synthetic Face",
        "repo": "prithivMLmods/Qwen-Image-Synthetic-Face",
        "weights": "qwen-synthetic-face.safetensors",
        "trigger_word": "Synthetic Face"
    },
    {
        "image": "https://huggingface.co/Tomechi02/Macne_style_enahncer/resolve/main/images/pixai-1913880604374308947-2.png",
        "title": "Macne Style Enahncer",
        "repo": "Tomechi02/Macne_style_enahncer",
        "weights": "Macne_Style_enhancer.safetensors",
        "trigger_word": "macloid, gomoku"
    },
    {
        "image": "https://huggingface.co/itspoidaman/qwenglitch/resolve/main/images/GyZTwJIbkAAhS4h.jpeg",
        "title": "Qwen Glitch",
        "repo": "itspoidaman/qwenglitch",
        "weights": "qwenglitch1.safetensors",
        "trigger_word": "qwenglitch"
    },
    {
        "image": "https://huggingface.co/alfredplpl/qwen-image-modern-anime-lora/resolve/main/sample1.jpg",
        "title": "Modern Anime Lora",
        "repo": "alfredplpl/qwen-image-modern-anime-lora",
        "weights": "lora.safetensors",
        "trigger_word": "Japanese modern anime style"
    },
    {
        "image": "https://huggingface.co/damnthatai/Apple_QuickTake_150_Digital_Camera_Qwen/resolve/main/images/20250817084713_Qwen.jpg",
        "title": "Apple QuickTake 150 Digital Camera",
        "repo": "damnthatai/Apple_QuickTake_150_Digital_Camera_Qwen",
        "weights": "quicktake150style_qwen.safetensors",
        "trigger_word": "quicktake150style"
    },
]

# Initialize the base model and autoencoders
dtype = torch.bfloat16
base_model = "Qwen/Qwen-Image"

# Initialize TAEF1 for fast previews and the standard VAE for high-quality final images
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)

# Scheduler configuration from the Qwen-Image-Lightning repository
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)

# Main pipeline for text-to-image, using taef1 for fast decoding during generation
pipe = DiffusionPipeline.from_pretrained(
    base_model, scheduler=scheduler, torch_dtype=dtype, vae=taef1
).to(device)

# Image-to-image pipeline, using the high-quality VAE
pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
    base_model,
    vae=good_vae,
    scheduler=scheduler,
    torch_dtype=dtype
).to(device)


# Lightning LoRA info (no global state)
LIGHTNING_LORA_REPO = "lightx2v/Qwen-Image-Lightning"
LIGHTNING_LORA_WEIGHT = "Qwen-Image-Lightning-8steps-V1.0.safetensors"

MAX_SEED = np.iinfo(np.int32).max

class Timer:
    def __init__(self, task_name=""):
        self.task_name = task_name

    def __enter__(self):
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        self.elapsed_time = self.end_time - self.start_time
        if self.task_name:
            print(f"Elapsed time for {self.task_name}: {self.elapsed_time:.6f} seconds")
        else:
            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")

def compute_image_dimensions(aspect_ratio):
    """Converts aspect ratio string to width, height tuple."""
    if aspect_ratio == "1:1":
        return 1024, 1024
    elif aspect_ratio == "16:9":
        return 1152, 640
    elif aspect_ratio == "9:16":
        return 640, 1152
    elif aspect_ratio == "4:3":
        return 1024, 768
    elif aspect_ratio == "3:4":
        return 768, 1024
    elif aspect_ratio == "3:2":
        return 1024, 688
    elif aspect_ratio == "2:3":
        return 688, 1024
    else:
        return 1024, 1024

def handle_lora_selection(evt: gr.SelectData, aspect_ratio):
    selected_lora = loras[evt.index]
    new_placeholder = f"Type a prompt for {selected_lora['title']}"
    lora_repo = selected_lora["repo"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"
    
    # Update aspect ratio if specified in LoRA config
    if "aspect" in selected_lora:
        if selected_lora["aspect"] == "portrait":
            aspect_ratio = "9:16"
        elif selected_lora["aspect"] == "landscape":
            aspect_ratio = "16:9"
        else:
            aspect_ratio = "1:1"
    
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        evt.index,
        aspect_ratio,
    )

def adjust_generation_mode(speed_mode):
    """Update UI based on speed/quality toggle."""
    if speed_mode == "Fast (8 steps)":
        return gr.update(value="Fast mode selected - 8 steps with Lightning LoRA"), 8, 1.0
    else: 
        return gr.update(value="Base mode selected - 48 steps for best quality"), 48, 4.0

def image_to_image_generation(prompt_mash, image_input, strength, steps, cfg_scale, width, height, lora_scale, seed):
    """Handles the image-to-image generation process."""
    generator = torch.Generator(device="cuda").manual_seed(seed)
    pipe_i2i.to("cuda")
    
    # Resize and convert input image
    image_input_pil = load_image(image_input).resize((width, height), Image.Resampling.LANCZOS)
    
    final_image = pipe_i2i(
        prompt=prompt_mash,
        image=image_input_pil,
        strength=strength,
        num_inference_steps=steps,
        guidance_scale=cfg_scale,
        generator=generator,
        # Note: image-to-image with Qwen doesn't use `true_cfg_scale`
    ).images[0]
    return final_image

@spaces.GPU(duration=100)
def process_generation_request(
    prompt, image_input, image_strength, cfg_scale, steps, selected_index, 
    randomize_seed, seed, aspect_ratio, lora_scale, speed_mode, progress=gr.Progress(track_tqdm=True)
):
    if selected_index is None:
        raise gr.Error("You must select a LoRA before proceeding.🧨")
    
    selected_lora = loras[selected_index]
    lora_path = selected_lora["repo"]
    trigger_word = selected_lora["trigger_word"]
    
    # Prepare prompt with trigger word
    if trigger_word:
        prompt_mash = f"{trigger_word}, {prompt}" if prompt else trigger_word
    else:
        prompt_mash = prompt

    # Set random seed if requested
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    # Determine which pipeline to use
    pipe_to_use = pipe_i2i if image_input is not None else pipe
    
    # Always unload any existing LoRAs first to avoid conflicts
    with Timer("Unloading existing LoRAs"):
        pipe_to_use.unload_lora_weights()

    # Load LoRAs based on speed mode
    if speed_mode == "Fast (8 steps)":
        with Timer("Loading Lightning LoRA and style LoRA"):
            pipe_to_use.load_lora_weights(
                LIGHTNING_LORA_REPO, 
                weight_name=LIGHTNING_LORA_WEIGHT,
                adapter_name="lightning"
            )
            weight_name = selected_lora.get("weights")
            pipe_to_use.load_lora_weights(
                lora_path, 
                weight_name=weight_name, 
                adapter_name="style"
            )
            pipe_to_use.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
    else: # Quality mode
        with Timer(f"Loading LoRA weights for {selected_lora['title']}"):
            weight_name = selected_lora.get("weights")
            pipe_to_use.load_lora_weights(lora_path, weight_name=weight_name)
    
    width, height = compute_image_dimensions(aspect_ratio)

    # --- Generation ---
    if image_input is not None:
        # Image-to-Image Generation
        final_image = image_to_image_generation(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
        yield final_image, seed, gr.update(visible=False)
    else:
        # Text-to-Image Generation with Previews
        pipe.to("cuda")
        generator = torch.Generator(device="cuda").manual_seed(seed)
        
        # Callback for generating previews
        def callback_on_step_end(pipe, step_index, timestep, callback_kwargs):
            latents = callback_kwargs["latents"]
            # Use the fast taef1 decoder for previews
            with torch.no_grad():
                image = pipe.decode_latents(latents.to(dtype))[0]
            progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_index + 1}; --total: {steps};"></div></div>'
            yield {"image": image, "seed": seed, "progress": gr.update(value=progress_bar, visible=True)}
            return callback_kwargs

        # Generate image with step-by-step previews
        with Timer("Generating image with previews"):
            generation_output = pipe(
                prompt=prompt_mash,
                num_inference_steps=steps,
                true_cfg_scale=cfg_scale,
                width=width,
                height=height,
                generator=generator,
                output_type="latent", # Get latents to decode with the good VAE later
                callback_on_step_end=callback_on_step_end
            )

        # Decode the final image with the high-quality VAE
        with Timer("Final decoding with good VAE"):
            final_latents = generation_output.images
            pipe.vae = good_vae # Temporarily swap to the good VAE
            final_image = pipe.decode_latents(final_latents.to(dtype))[0]
            pipe.vae = taef1 # Swap back to taef1 for the next run
            
        yield final_image, seed, gr.update(visible=False)


def fetch_hf_adapter_files(link):
    split_link = link.split("/")
    if len(split_link) != 2:
        raise Exception("Invalid Hugging Face repository link format.")

    print(f"Repository attempted: {split_link}")
    
    model_card = ModelCard.load(link)
    base_model = model_card.data.get("base_model")
    print(f"Base model: {base_model}")

    acceptable_models = {"Qwen/Qwen-Image"}
    models_to_check = base_model if isinstance(base_model, list) else [base_model]
    
    if not any(model in acceptable_models for model in models_to_check):
        raise Exception("Not a Qwen-Image LoRA!")
        
    image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url")
    trigger_word = model_card.data.get("instance_prompt", "")
    image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None

    fs = HfFileSystem()
    try:
        list_of_files = fs.ls(link, detail=False)
        safetensors_name = next((f.split('/')[-1] for f in list_of_files if f.endswith(".safetensors")), None)
        if not safetensors_name:
            raise Exception("No valid *.safetensors file found in the repository.")
    except Exception as e:
        print(e)
        raise Exception("Could not find a valid *.safetensors file in the Hugging Face repository.")
    
    return split_link[1], link, safetensors_name, trigger_word, image_url

def validate_custom_adapter(link):
    print(f"Checking a custom model on: {link}")
    if link.startswith("https://huggingface.co"):
        link = urlparse(link).path.strip("/")
    return fetch_hf_adapter_files(link)

def incorporate_custom_adapter(custom_lora):
    global loras
    if custom_lora:
        try:
            title, repo, path, trigger_word, image = validate_custom_adapter(custom_lora)
            print(f"Loaded custom LoRA: {repo}")
            card = f'''
            <div class="custom_lora_card">
              <span>Loaded custom LoRA:</span>
              <div class="card_internal">
                <img src="{image}" />
                <div>
                    <h3>{title}</h3>
                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
                </div>
              </div>
            </div>
            '''
            existing_item_index = next((i for i, item in enumerate(loras) if item['repo'] == repo), None)
            if existing_item_index is None:
                new_item = {"image": image, "title": title, "repo": repo, "weights": path, "trigger_word": trigger_word}
                loras.append(new_item)
                existing_item_index = len(loras) - 1
        
            return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
        except Exception as e:
            gr.Warning(f"Invalid LoRA: {e}")
            return gr.update(visible=True, value=f"Invalid LoRA: {e}"), gr.update(visible=True), gr.update(), "", None, ""
    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""

def discard_custom_adapter():
    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""


css = '''
#gen_btn{height: 100%}
#gen_column{align-self: stretch}
#title{text-align: center}
#title h1{font-size: 3em; display:inline-flex; align-items:center}
#title img{width: 100px; margin-right: 0.5em}
#gallery .grid-wrap{height: 10vh}
#lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
.card_internal{display: flex;height: 100px;margin-top: .5em}
.card_internal img{margin-right: 1em}
.styler{--form-gap-width: 0px !important}
#speed_status{padding: .5em; border-radius: 5px; margin: 1em 0}
#progress{height:30px}
#progress .generating{display:none}
.progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
.progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.1s ease-in-out}
'''

with gr.Blocks(theme="bethecloud/storj_theme", css=css, delete_cache=(120, 120)) as app:
    title = gr.HTML("""<h1>Qwen Image LoRA DLC⛵</h1>""", elem_id="title")
    selected_index = gr.State(None)
    
    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="✦︎ Choose the LoRA and type the prompt")
        with gr.Column(scale=1, elem_id="gen_column"):
            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
    
    with gr.Row():
        with gr.Column():
            selected_info = gr.Markdown("")
            gallery = gr.Gallery(
                [(item["image"], item["title"]) for item in loras],
                label="LoRA Gallery",
                allow_preview=False,
                columns=3,
                elem_id="gallery",
                show_share_button=False
            )
            with gr.Accordion("Image-to-Image (Optional)", open=False):
                image_input = gr.Image(type="filepath", label="Input Image")
                image_strength = gr.Slider(label="Image Strength", minimum=0.1, maximum=1.0, step=0.05, value=0.6)

            with gr.Group():
                custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path", placeholder="username/lora-model-name")
                gr.Markdown("[Check Qwen-Image LoRAs](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image)", elem_id="lora_list")
            custom_lora_info = gr.HTML(visible=False)
            custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
        
        with gr.Column():
            result = gr.Image(label="Generated Image")
            progress_bar = gr.HTML(visible=False, elem_id="progress")

            with gr.Row():
                aspect_ratio = gr.Dropdown(
                    label="Aspect Ratio",
                    choices=["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"],
                    value="1:1"
                )
            with gr.Row():
                speed_mode = gr.Dropdown(
                    label="Output Mode",
                    choices=["Fast (8 steps)", "Base (48 steps)"],
                    value="Base (48 steps)",
                )
            
            speed_status = gr.Markdown("Base mode selected", elem_id="speed_status")

    with gr.Row():
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Column():
                with gr.Row():
                    cfg_scale = gr.Slider(
                        label="Guidance Scale", 
                        minimum=1.0, 
                        maximum=5.0, 
                        step=0.1, 
                        value=4.0,
                        info="Lower for speed mode, higher for quality. Also called 'True CFG'."
                    )
                    steps = gr.Slider(
                        label="Steps", 
                        minimum=4, 
                        maximum=50, 
                        step=1, 
                        value=48,
                        info="Automatically set by speed mode"
                    )
                
                with gr.Row():
                    randomize_seed = gr.Checkbox(True, label="Randomize seed")
                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=2, step=0.01, value=1.0)

    # Event handlers
    gallery.select(
        handle_lora_selection,
        inputs=[aspect_ratio],
        outputs=[prompt, selected_info, selected_index, aspect_ratio]
    )
    
    speed_mode.change(
        adjust_generation_mode,
        inputs=[speed_mode],
        outputs=[speed_status, steps, cfg_scale]
    )
    
    custom_lora.input(
        incorporate_custom_adapter,
        inputs=[custom_lora],
        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
    )
    
    custom_lora_button.click(
        discard_custom_adapter,
        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
    )
    
    gen_inputs = [prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, aspect_ratio, lora_scale, speed_mode]
    gen_outputs = [result, seed, progress_bar]
    
    generate_button.click(
        fn=process_generation_request,
        inputs=gen_inputs,
        outputs=gen_outputs
    )
    prompt.submit(
        fn=process_generation_request,
        inputs=gen_inputs,
        outputs=gen_outputs
    )

app.queue()
app.launch(share=False, ssr_mode=False, show_error=True)