Update app.py

app.py

@@ -1,106 +1,220 @@
-import math
+import os
 import gradio as gr
-import spaces
+import json
+import logging
+import torch
 from PIL import Image
+import spaces
+from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
+import copy
+import random
+import time
+import re
+import math
+import numpy as np
+import traceback
+
+#from diffsynth_engine import fetch_model, QwenImagePipeline, QwenImagePipelineConfig
 
-from diffsynth_engine import (
-    fetch_model,
-    QwenImagePipeline,
-    QwenImagePipelineConfig,
-)
-
-# ----------------------------
-# Pipeline initialization (CPU)
-# ----------------------------
-config = QwenImagePipelineConfig.basic_config(
-    model_path=fetch_model(
-        "Qwen/Qwen-Image-2512",
-        path="transformer/*.safetensors"
-    ),
-    encoder_path=fetch_model(
-        "Qwen/Qwen-Image-2512",
-        path="text_encoder/*.safetensors"
-    ),
-    vae_path=fetch_model(
-        "Qwen/Qwen-Image-2512",
-        path="vae/*.safetensors"
-    ),
-    offload_mode="cpu_offload",
-)
-
-pipe = QwenImagePipeline.from_pretrained(config)
-
-# Load Turbo LoRA
-pipe.load_lora(
-    path=fetch_model(
-        "Wuli-Art/Qwen-Image-2512-Turbo-LoRA",
-        path="Wuli-Qwen-Image-2512-Turbo-LoRA-4steps-V1.0-bf16.safetensors"
-    ),
-    scale=1.0,
-    fused=True,
-)
-
-# Scheduler config
-pipe.apply_scheduler_config(
-    {
-        "exponential_shift_mu": math.log(2.5),
-        "use_dynamic_shifting": True,
-        "shift_terminal": None,
+from huggingface_hub import login, InferenceClient
+
+login(token=os.environ.get('hf'))
+
+def apply_aspect_ratio(ratio):
+    sizes = {
+        "1:1": (1024, 1024),
+        "16:9": (1365, 768),
+        "9:16": (768, 1365),
+        "3:2": (1254, 836),
+        "2:3": (836, 1254),
+        "3:1": (1774, 591),
+        "2:1": (1448, 724),
     }
-)
-
-# ----------------------------
-# GPU inference (ZeroGPU)
-# ----------------------------
-@spaces.GPU
-def generate(
-    prompt: str,
-    seed: int,
-    width: int,
-    height: int,
+    return sizes.get(ratio, (1024, 1024))
+
+
+# Initialize the base model
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+base_model = "Qwen/Qwen-Image-2512"
+
+# 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)
+pipe = DiffusionPipeline.from_pretrained(
+    base_model, scheduler=scheduler, torch_dtype=dtype
+).to(device)
+
+
+
+# Lightning LoRA info (no global state)
+LIGHTNING_LORA_REPO = "lightx2v/Qwen-Image-2512-Lightning"
+LIGHTNING_LORA_WEIGHT = "Qwen-Image-2512-Lightning-4steps-V1.0-bf16.safetensors"
+LIGHTNING8_LORA_WEIGHT = "Qwen-Image-Lightning-8steps-V2.0-bf16.safetensors"
+LIGHTNING_FP8_4STEPS_LORA_WEIGHT = "Qwen-Image-fp8-e4m3fn-Lightning-4steps-V1.0-bf16.safetensors"
+
+#LIGHTNING_LORA_REPO = "Wuli-art/Qwen-Image-2512-Turbo-LoRA"
+#LIGHTNING_LORA_WEIGHT = "Wuli-Qwen-Image-2512-Turbo-LoRA-4steps-V1.0-bf16.safetensors"
+#LIGHTNING8_LORA_WEIGHT = "Wuli-Qwen-Image-2512-Turbo-LoRA-4steps-V1.0-bf16.safetensors"
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_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.activity_name:
+            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
+        else:
+            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
+
+
+@spaces.GPU(duration=70)
+def generate_image(
+    prompt_mash,
+    width,
+    height,
 ):
-    result = pipe(
-        prompt=prompt,
-        cfg_scale=1,
-        num_inference_steps=4,
-        seed=seed,
-        width=width,
-        height=height,
-    )
+    pipe.to("cuda")
+    if negative_prompt == '':
+        negative_prompt =  "低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。"
 
-    # result is PIL.Image
-    return result
 
+    seed = 235234
+    num_images = 1
+    seeds = [seed + (i * 100) for i in range(num_images)]
+    generators = [torch.Generator(device="cuda").manual_seed(s) for s in seeds]
+    
+    images = []
 
-# ----------------------------
-# Gradio UI
-# ----------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("## Qwen Image · ZeroGPU demo")
+    with calculateDuration("Generating images (sequential)"):
+        for i in range(num_images):
+            current_seed = seed + (i * 100)
+            generator = torch.Generator(device="cuda").manual_seed(current_seed)
+    
+            result = pipe(
+                prompt=prompt_mash,
+                negative_prompt=negative_prompt,
+                num_inference_steps=4,
+                true_cfg_scale=1,
+                width=width,
+                height=height,
+                num_images_per_prompt=1,
+                generator=generator,
+            )
+    
+            images.append((result.images[0], current_seed))
+    return images
+           
 
-    with gr.Row():
-        prompt = gr.Textbox(
-            label="Prompt",
-            value="a young girl with flowing long hair, wearing a white halter dress and smiling sweetly. The background features a blue seaside where seagulls fly freely.",
-            lines=4,
+
+@spaces.GPU(duration=70)
+def run_lora(
+    prompt, 
+    width, 
+    height, 
+    progress=gr.Progress(track_tqdm=True)
+):
+
+    with calculateDuration("Loading Lightning LoRA and style LoRA"):
+        pipe.load_lora_weights(
+            'Wuli-Art/Qwen-Image-2512-Turbo-LoRA', 
+            weight_name='Wuli-Qwen-Image-2512-Turbo-LoRA-4steps-V1.0-bf16.safetensors',
+            adapter_name="lightning"
         )
+        pipe.set_adapters(["lightning"], adapter_weights=[1.0])
 
-    with gr.Row():
-        seed = gr.Number(value=42, label="Seed", precision=0)
-        width = gr.Number(value=1328, label="Width", precision=0)
-        height = gr.Number(value=1328, label="Height", precision=0)
+    
+    multiplier = float(quality_multiplier.replace('x', ''))
+    width = int(width * multiplier)
+    height = int(height * multiplier)
+    num_images = int(quantity) + 1
 
-    generate_btn = gr.Button("Generate")
 
-    output_image = gr.Image(
-        label="Generated image",
-        type="pil",
+    pairs = generate_image(
+        prompt,
+        width,
+        height
     )
 
-    generate_btn.click(
-        fn=generate,
-        inputs=[prompt, seed, width, height],
-        outputs=output_image,
+    images_for_gallery = [
+        (img, str(s))
+        for (img, s) in pairs
+    ]
+
+    return images_for_gallery
+
+
+
+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}
+'''
+
+
+with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 60)) as app:
+    title = gr.HTML(
+        """<img src=\"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_logo.png\" alt=\"Qwen-Image\" style=\"width: 280px; margin: 0 auto\">        
+        <h3 style=\"margin-top: -10px\">Wuli-art/Qwen-Image-2512-Turbo-LoRA</h3>""",
+        elem_id="title",
     )
+    
+    selected_index = gr.State(None)
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Type a prompt after selecting a LoRA")
+            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn", interactive=False)
+        with gr.Column(scale=1, elem_id="gen_column"):
+            result = gr.Gallery(label="Generated Images", show_label=True, elem_id="result_gallery")
+    
+
+    generate_event = gr.on(
+        triggers=[generate_button.click, prompt.submit],
+        fn=run_lora,
+        inputs=[
+            prompt, 
+            width, 
+            height
+        ],
+        outputs=[result]
+    )
+
 
-demo.queue().launch()
+app.queue()
+app.launch()