prithivMLmods's picture
Update app.py
151bae0 verified
raw
history blame
21.3 kB
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,
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
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:", processing_device)
# List of predefined style models (formerly LoRAs)
style_definitions = [
{
"thumbnail_url": "https://huggingface.co/prithivMLmods/Qwen-Image-Studio-Realism/resolve/main/images/2.png",
"style_name": "Studio Realism",
"repo_id": "prithivMLmods/Qwen-Image-Studio-Realism",
"weight_file": "qwen-studio-realism.safetensors",
"activation_phrase": "Studio Realism"
},
{
"thumbnail_url": "https://huggingface.co/prithivMLmods/Qwen-Image-Sketch-Smudge/resolve/main/images/1.png",
"style_name": "Sketch Smudge",
"repo_id": "prithivMLmods/Qwen-Image-Sketch-Smudge",
"weight_file": "qwen-sketch-smudge.safetensors",
"activation_phrase": "Sketch Smudge"
},
{
"thumbnail_url": "https://huggingface.co/prithivMLmods/Qwen-Image-Anime-LoRA/resolve/main/images/1.png",
"style_name": "Qwen Anime",
"repo_id": "prithivMLmods/Qwen-Image-Anime-LoRA",
"weight_file": "qwen-anime.safetensors",
"activation_phrase": "Qwen Anime"
},
{
"thumbnail_url": "https://huggingface.co/prithivMLmods/Qwen-Image-Synthetic-Face/resolve/main/images/2.png",
"style_name": "Synthetic Face",
"repo_id": "prithivMLmods/Qwen-Image-Synthetic-Face",
"weight_file": "qwen-synthetic-face.safetensors",
"activation_phrase": "Synthetic Face"
},
{
"thumbnail_url": "https://huggingface.co/prithivMLmods/Qwen-Image-Fragmented-Portraiture/resolve/main/images/3.png",
"style_name": "Fragmented Portraiture",
"repo_id": "prithivMLmods/Qwen-Image-Fragmented-Portraiture",
"weight_file": "qwen-fragmented-portraiture.safetensors",
"activation_phrase": "Fragmented Portraiture"
},
]
# --- Model Initialization ---
model_precision = torch.bfloat16
processing_device = "cuda" if torch.cuda.is_available() else "cpu"
foundation_model_id = "Qwen/Qwen-Image"
# Sampler configuration from the Qwen-Image-Lightning repository
sampler_settings = {
"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,
}
sampler = FlowMatchEulerDiscreteScheduler.from_config(sampler_settings)
diffusion_pipeline = DiffusionPipeline.from_pretrained(
foundation_model_id, scheduler=sampler, torch_dtype=model_precision
).to(processing_device)
# Information for the fast generation LoRA
FAST_GENERATION_LORA_REPO = "lightx2v/Qwen-Image-Lightning"
FAST_GENERATION_LORA_WEIGHTS = "Qwen-Image-Lightning-8steps-V1.0.safetensors"
MAX_SEED_VALUE = np.iinfo(np.int32).max
class ExecutionTimer:
"""A context manager to time a block of code."""
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
activity_log = f" for {self.activity_name}" if self.activity_name else ""
print(f"Elapsed time{activity_log}: {self.elapsed_time:.6f} seconds")
def get_dimensions_from_ratio(aspect_ratio_str):
"""Converts an aspect ratio string to a (width, height) tuple."""
ratios = {
"1:1": (1024, 1024),
"16:9": (1152, 640),
"9:16": (640, 1152),
"4:3": (1024, 768),
"3:4": (768, 1024),
"3:2": (1024, 688),
"2:3": (688, 1024),
}
return ratios.get(aspect_ratio_str, (1024, 1024))
def on_style_select(event_data: gr.SelectData, current_aspect_ratio):
"""Handles the user selecting a style from the gallery."""
selected_style = style_definitions[event_data.index]
new_placeholder = f"Type a prompt for {selected_style['style_name']}"
repo_id = selected_style["repo_id"]
updated_info_text = f"### Selected: [{repo_id}](https://huggingface.co/{repo_id}) ✨"
# Update aspect ratio if specified in the style's configuration
if "aspect" in selected_style:
if selected_style["aspect"] == "portrait":
current_aspect_ratio = "9:16"
elif selected_style["aspect"] == "landscape":
current_aspect_ratio = "16:9"
else:
current_aspect_ratio = "1:1"
return (
gr.update(placeholder=new_placeholder),
updated_info_text,
event_data.index,
current_aspect_ratio,
)
def on_mode_change(generation_mode):
"""Updates UI elements based on the selected generation mode (Speed/Quality)."""
if generation_mode == "Speed (8 steps)":
return gr.update(value="Speed mode selected - 8 steps with Lightning LoRA"), 8, 1.0
else:
return gr.update(value="Quality mode selected - 45 steps for best quality"), 45, 3.5
@spaces.GPU(duration=70)
def execute_image_generation(full_prompt, steps, seed_val, cfg, width, height, negative_prompt=""):
"""Generates an image using the diffusion pipeline."""
diffusion_pipeline.to("cuda")
generator = torch.Generator(device="cuda").manual_seed(seed_val)
with ExecutionTimer("Image Generation"):
generated_image = diffusion_pipeline(
prompt=full_prompt,
negative_prompt=negative_prompt,
num_inference_steps=steps,
true_cfg_scale=cfg,
width=width,
height=height,
generator=generator,
).images[0]
return generated_image
@spaces.GPU(duration=70)
def handle_generate_request(prompt_text, cfg, steps, style_idx, use_random_seed, seed_val, aspect_ratio_str, style_scale, generation_mode, progress=gr.Progress(track_tqdm=True)):
"""Main function to handle a user's image generation request."""
if style_idx is None:
raise gr.Error("You must select a style before generating an image.")
selected_style = style_definitions[style_idx]
style_repo_path = selected_style["repo_id"]
activation_phrase = selected_style["activation_phrase"]
# Combine the user prompt with the style's activation phrase
if activation_phrase:
position = selected_style.get("trigger_position", "prepend")
if position == "prepend":
full_prompt = f"{activation_phrase} {prompt_text}"
else:
full_prompt = f"{prompt_text} {activation_phrase}"
else:
full_prompt = prompt_text
# Always unload existing adapters to start fresh
with ExecutionTimer("Unloading existing adapters"):
diffusion_pipeline.unload_lora_weights()
# Load adapters based on the selected generation mode
if generation_mode == "Speed (8 steps)":
with ExecutionTimer("Loading Lightning and Style adapters"):
# Load the fast generation adapter first
diffusion_pipeline.load_lora_weights(
FAST_GENERATION_LORA_REPO,
weight_name=FAST_GENERATION_LORA_WEIGHTS,
adapter_name="lightning"
)
# Load the selected style adapter
weight_file = selected_style.get("weight_file", None)
diffusion_pipeline.load_lora_weights(
style_repo_path,
weight_name=weight_file,
low_cpu_mem_usage=True,
adapter_name="style"
)
# Set both adapters active with their respective weights
diffusion_pipeline.set_adapters(["lightning", "style"], adapter_weights=[1.0, style_scale])
else: # Quality mode
with ExecutionTimer(f"Loading adapter weights for {selected_style['style_name']}"):
weight_file = selected_style.get("weight_file", None)
diffusion_pipeline.load_lora_weights(
style_repo_path,
weight_name=weight_file,
low_cpu_mem_usage=True
)
# Set the seed for reproducibility
with ExecutionTimer("Setting seed"):
if use_random_seed:
seed_val = random.randint(0, MAX_SEED_VALUE)
# Get image dimensions
width, height = get_dimensions_from_ratio(aspect_ratio_str)
# Generate the final image
final_image = execute_image_generation(full_prompt, steps, seed_val, cfg, width, height)
return final_image, seed_val
def fetch_hf_safetensors_details(repo_link):
"""Fetches details of a LoRA from a Hugging Face repository."""
split_link = repo_link.split("/")
if len(split_link) != 2:
raise ValueError("Invalid Hugging Face repository link format.")
print(f"Attempting to load repository: {repo_link}")
model_card = ModelCard.load(repo_link)
base_model = model_card.data.get("base_model")
print(f"Base model identified: {base_model}")
# Validate that the LoRA is compatible with Qwen-Image
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 TypeError("The provided model is not a Qwen-Image compatible LoRA.")
# Extract metadata from the model card
image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url")
activation_phrase = model_card.data.get("instance_prompt", "")
image_url = f"https://huggingface.co/{repo_link}/resolve/main/{image_path}" if image_path else None
# Find the .safetensors file in the repository
fs = HfFileSystem()
try:
repo_files = fs.ls(repo_link, detail=False)
safetensors_filename = None
for file_path in repo_files:
filename = file_path.split("/")[-1]
if filename.endswith(".safetensors"):
safetensors_filename = filename
break
if not safetensors_filename:
raise FileNotFoundError("No .safetensors file was found in the repository.")
except Exception as e:
print(e)
raise IOError("Could not access the Hugging Face repository or find a valid .safetensors file.")
return split_link[1], repo_link, safetensors_filename, activation_phrase, image_url
def parse_custom_model_source(source_text):
"""Parses a user-provided link to a custom LoRA."""
print(f"Parsing custom model source: {source_text}")
if source_text.endswith('.safetensors') and 'huggingface.co' in source_text:
parts = source_text.split('/')
try:
hf_index = parts.index('huggingface.co')
username = parts[hf_index + 1]
repo_name = parts[hf_index + 2]
repo_id = f"{username}/{repo_name}"
safetensors_filename = parts[-1]
try:
model_card = ModelCard.load(repo_id)
activation_phrase = model_card.data.get("instance_prompt", "")
image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url")
image_url = f"https://huggingface.co/{repo_id}/resolve/main/{image_path}" if image_path else None
except Exception:
activation_phrase = ""
image_url = None
return repo_name, repo_id, safetensors_filename, activation_phrase, image_url
except ValueError:
raise ValueError("Invalid .safetensors URL format.")
if source_text.startswith("https://"):
parsed_url = urlparse(source_text)
if "huggingface.co" in parsed_url.netloc:
repo_link = parsed_url.path.strip("/")
return fetch_hf_safetensors_details(repo_link)
# Assume it's a direct repo path like "username/repo-name"
return fetch_hf_safetensors_details(source_text)
def add_custom_style_model(custom_model_path):
"""Adds a custom LoRA provided by the user to the session."""
global style_definitions
if custom_model_path:
try:
style_name, repo_id, weight_file, activation_phrase, thumbnail_url = parse_custom_model_source(custom_model_path)
print(f"Successfully loaded custom style: {repo_id}")
card_html = f'''
<div class="custom_lora_card">
<span>Loaded custom style:</span>
<div class="card_internal">
<img src="{thumbnail_url}" alt="{style_name}" />
<div>
<h3>{style_name}</h3>
<small>{"Activation phrase: <code><b>"+activation_phrase+"</b></code>" if activation_phrase else "No activation phrase found. If required, include it in your prompt."}<br></small>
</div>
</div>
</div>
'''
# Check if this style already exists
existing_item_index = next((index for (index, item) in enumerate(style_definitions) if item['repo_id'] == repo_id), None)
if existing_item_index is None:
new_style_item = {
"thumbnail_url": thumbnail_url,
"style_name": style_name,
"repo_id": repo_id,
"weight_file": weight_file,
"activation_phrase": activation_phrase
}
style_definitions.append(new_style_item)
existing_item_index = len(style_definitions) - 1
return gr.update(visible=True, value=card_html), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {weight_file}", existing_item_index, activation_phrase
except Exception as e:
gr.Warning(f"Failed to load custom style. Error: {e}")
error_message = f"Invalid input. Could not load the specified style. Please check the link or repository path."
return gr.update(visible=True, value=error_message), gr.update(visible=True), gr.update(), "", None, ""
# If input is empty, hide the custom section
return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
def remove_custom_style_model():
"""Resets the UI when a custom LoRA is removed."""
return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
# --- Gradio UI Definition ---
app_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; object-fit: cover;}
.styler{--form-gap-width: 0px !important}
#speed_status{padding: .5em; border-radius: 5px; margin: 1em 0}
.custom_lora_card{padding: 1em; border: 1px solid var(--border-color-primary); border-radius: var(--radius-lg)}
'''
with gr.Blocks(theme="bethecloud/storj_theme", css=app_css, delete_cache=(120, 120)) as web_interface:
main_title = gr.HTML("""<h1>Qwen Image LoRA DLC❤️‍🔥</h1>""", elem_id="title")
selected_style_index = gr.State(None)
with gr.Row():
with gr.Column(scale=3):
prompt_textbox = gr.Textbox(label="Prompt", lines=1, placeholder="Select a style to begin...")
with gr.Column(scale=1, elem_id="gen_column"):
generate_btn = gr.Button("Generate", variant="primary", elem_id="gen_btn")
with gr.Row():
with gr.Column():
selected_style_info = gr.Markdown("")
style_gallery = gr.Gallery(
[(item["thumbnail_url"], item["style_name"]) for item in style_definitions],
label="Style Gallery",
allow_preview=False,
columns=3,
elem_id="gallery",
show_share_button=False
)
with gr.Group():
custom_style_textbox = gr.Textbox(label="Load Custom Style", info="Enter a Hugging Face repository path (e.g., username/repo-name)", placeholder="username/qwen-image-custom-style")
gr.Markdown("[Find More Qwen-Image Styles Here](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image)", elem_id="lora_list")
custom_style_info_html = gr.HTML(visible=False)
remove_custom_style_btn = gr.Button("Remove Custom Style", visible=False)
with gr.Column():
output_image_display = gr.Image(label="Generated Image")
with gr.Row():
aspect_ratio_dropdown = 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():
generation_mode_dropdown = gr.Dropdown(
label="Generation Mode",
choices=["Speed (8 steps)", "Quality (45 steps)"],
value="Quality (45 steps)",
)
generation_mode_status_display = gr.Markdown("Quality mode active", elem_id="speed_status")
with gr.Row():
with gr.Accordion("Advanced Settings", open=False):
with gr.Column():
with gr.Row():
cfg_scale_slider = gr.Slider(
label="Guidance Scale (CFG)",
minimum=1.0,
maximum=5.0,
step=0.1,
value=3.5,
info="Adjusts how strictly the model follows the prompt. Lower for speed, higher for quality."
)
steps_slider = gr.Slider(
label="Inference Steps",
minimum=4,
maximum=50,
step=1,
value=45,
info="Number of steps for the generation process. Automatically set by Generation Mode."
)
with gr.Row():
randomize_seed_checkbox = gr.Checkbox(True, label="Use Random Seed")
seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED_VALUE, step=1, value=0, randomize=True)
style_scale_slider = gr.Slider(label="Style Strength", minimum=0, maximum=2, step=0.01, value=1.0)
# --- Event Handlers ---
style_gallery.select(
on_style_select,
inputs=[aspect_ratio_dropdown],
outputs=[prompt_textbox, selected_style_info, selected_style_index, aspect_ratio_dropdown]
)
generation_mode_dropdown.change(
on_mode_change,
inputs=[generation_mode_dropdown],
outputs=[generation_mode_status_display, steps_slider, cfg_scale_slider]
)
custom_style_textbox.submit(
add_custom_style_model,
inputs=[custom_style_textbox],
outputs=[custom_style_info_html, remove_custom_style_btn, style_gallery, selected_style_info, selected_style_index, prompt_textbox]
)
remove_custom_style_btn.click(
remove_custom_style_model,
outputs=[custom_style_info_html, remove_custom_style_btn, style_gallery, selected_style_info, selected_style_index, custom_style_textbox]
)
# Combined trigger for generation
generate_triggers = [generate_btn.click, prompt_textbox.submit]
gr.on(
triggers=generate_triggers,
fn=handle_generate_request,
inputs=[prompt_textbox, cfg_scale_slider, steps_slider, selected_style_index, randomize_seed_checkbox, seed_slider, aspect_ratio_dropdown, style_scale_slider, generation_mode_dropdown],
outputs=[output_image_display, seed_slider]
)
web_interface.queue()
web_interface.launch(share=False, ssr_mode=False, show_error=True)