import os import json import math import time import random import numpy as np from typing import Iterable import torch from PIL import Image import gradio as gr import spaces from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler from huggingface_hub import HfFileSystem, ModelCard from gradio.themes import Soft from gradio.themes.utils import colors, fonts, sizes # ── THEME ──────────────────────────────────────────────────────────────────── THEME_PRESETS = { "orange-red": {"primary": "#c2410c", "secondary": "#ea580c"}, "violet": {"primary": "#7c3aed", "secondary": "#8b5cf6"}, "ocean": {"primary": "#0369a1", "secondary": "#0ea5e9"}, "emerald": {"primary": "#226105", "secondary": "#73ff00"}, "rose": {"primary": "#be185d", "secondary": "#ec4899"}, "amber": {"primary": "#b45309", "secondary": "#f59e0b"}, "cyan": {"primary": "#0e7490", "secondary": "#06b6d4"}, } def make_color(name: str, hex_val: str) -> colors.Color: """Generate a Gradio Color object from a single hex value (simplified ramp).""" return colors.Color( name=name, c50=hex_val + "15", c100=hex_val + "25", c200=hex_val + "40", c300=hex_val + "60", c400=hex_val + "80", c500=hex_val, c600=hex_val, c700=hex_val, c800=hex_val, c900=hex_val, c950=hex_val, ) class DRexTheme(Soft): def __init__( self, primary_hue=colors.gray, accent_hex: str = "#c2410c", *, neutral_hue=colors.slate, text_size=sizes.text_md, font: Iterable[fonts.Font | str] = ( fonts.GoogleFont("Space Grotesk"), "Arial", "sans-serif" ), font_mono: Iterable[fonts.Font | str] = ( fonts.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace" ), ): super().__init__( primary_hue=primary_hue, neutral_hue=neutral_hue, text_size=text_size, font=font, font_mono=font_mono, ) a = accent_hex super().set( # backgrounds background_fill_primary="*primary_50", background_fill_primary_dark="*primary_900", body_background_fill="#0f0d0c", body_background_fill_dark="#0b0a09", # buttons button_primary_text_color="white", button_primary_text_color_hover="white", button_primary_background_fill=a, button_primary_background_fill_hover=a, button_primary_background_fill_dark=a, button_primary_background_fill_hover_dark=a, button_secondary_background_fill="*primary_200", button_secondary_background_fill_hover="*primary_300", # slider slider_color=a, slider_color_dark=a, # block block_title_text_weight="600", block_border_width="1px", block_shadow="none", block_label_background_fill="*primary_100", # input input_background_fill="#1a1816", input_background_fill_dark="#1a1816", input_border_color="#272422", input_border_color_dark="#272422", ) drex_theme = DRexTheme(accent_hex="#c2410c") # ── GPU / DEVICE ───────────────────────────────────────────────────────────── device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.bfloat16 print(f"[D-REX] Using device: {device}") # ── LORAS ──────────────────────────────────────────────────────────────────── loras = [ { "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/Shakker-Labs/AWPortrait-QW/resolve/main/images/08fdaf6b644b61136340d5c908ca37993e47f34cdbe2e8e8251c4c72.jpg", "title": "AWPortrait QW", "repo": "Shakker-Labs/AWPortrait-QW", "weights": "AWPortrait-QW_1.0.safetensors", "trigger_word": "Portrait", }, { "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/flymy-ai/qwen-image-realism-lora/resolve/main/assets/flymy_realism.png", "title": "Image Realism", "repo": "flymy-ai/qwen-image-realism-lora", "weights": "flymy_realism.safetensors", "trigger_word": "Super Realism Portrait", }, { "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/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", "repo": "alfredplpl/qwen-image-modern-anime-lora", "weights": "lora.safetensors", "trigger_word": "Japanese modern anime style", }, ] # ── MODEL INIT ──────────────────────────────────────────────────────────────── 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) # ZeroGPU: load on CPU — moved to cuda inside @spaces.GPU function only pipe = DiffusionPipeline.from_pretrained( "Qwen/Qwen-Image", scheduler=scheduler, torch_dtype=dtype ) LIGHTNING_REPO = "lightx2v/Qwen-Image-Lightning" LIGHTNING_WEIGHT = "Qwen-Image-Lightning-8steps-V1.0.safetensors" MAX_SEED = np.iinfo(np.int32).max ENHANCE_SUFFIXES = [ ", ultra detailed, cinematic lighting, 8k resolution, professional photography", ", masterpiece, intricate details, dramatic composition, volumetric light", ", photorealistic, sharp focus, studio lighting, high contrast, award winning", ", concept art, highly detailed, trending on artstation, vivid colors, epic", ", hyperrealistic, golden hour lighting, bokeh, atmospheric depth, stunning", ] # ── HELPERS ─────────────────────────────────────────────────────────────────── def aspect_to_wh(aspect: str): mapping = { "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 mapping.get(aspect, (1024, 1024)) def build_metadata_html(model_title: str, seed: int, steps: int, cfg: float, aspect: str) -> str: w, h = aspect_to_wh(aspect) return f"""
{model_title.split()[0]}
MODEL
{seed}
SEED
{steps}
STEPS
{cfg:.1f}
CFG
{w}×{h}
SIZE
""" def format_history_html(history: list) -> str: if not history: return "
no history yet
" rows = "" for h in reversed(history[-10:]): rows += f"""
{h['prompt']}
{h['model']} · {h['time']}
""" return f"
{rows}
" # ── CORE LOGIC ──────────────────────────────────────────────────────────────── @spaces.GPU(duration=200) def generate( prompt, neg_prompt, cfg, steps, selected_index, randomize_seed, seed, aspect, lora_scale, speed_mode, history_state, oauth_token: gr.OAuthToken | None = None, progress=gr.Progress(track_tqdm=True), ): if oauth_token is None: raise gr.Error("Please sign in with your HuggingFace account to generate images.") if selected_index is None: raise gr.Error("Select a LoRA from the gallery first.") if not prompt.strip(): raise gr.Error("Write a prompt before generating.") lora = loras[selected_index] trigger = lora["trigger_word"] prompt_in = f"{trigger} {prompt}" if trigger else prompt # ZeroGPU: LoRA loading must happen inside GPU scope pipe.to("cuda") try: pipe.unload_lora_weights() if speed_mode == "Fast · 8 steps": pipe.load_lora_weights(LIGHTNING_REPO, weight_name=LIGHTNING_WEIGHT, adapter_name="lightning") pipe.load_lora_weights(lora["repo"], weight_name=lora.get("weights"), low_cpu_mem_usage=True, adapter_name="style") pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale]) else: pipe.load_lora_weights(lora["repo"], weight_name=lora.get("weights"), low_cpu_mem_usage=True, adapter_name="style") pipe.set_adapters(["style"], adapter_weights=[lora_scale]) if randomize_seed: seed = random.randint(0, MAX_SEED) w, h = aspect_to_wh(aspect) generator = torch.Generator(device="cuda").manual_seed(seed) image = pipe( prompt=prompt_in, negative_prompt=neg_prompt, num_inference_steps=steps, true_cfg_scale=cfg, width=w, height=h, generator=generator, ).images[0] finally: pipe.to("cpu") torch.cuda.empty_cache() # update history history_state = history_state or [] history_state.append({ "prompt": prompt[:80], "model": lora["title"], "time": time.strftime("%H:%M"), }) history_state = history_state[-20:] meta_html = build_metadata_html(lora["title"], seed, steps, cfg, aspect) history_html = format_history_html(history_state) return image, seed, meta_html, history_html, history_state def enhance_prompt(prompt: str) -> str: if not prompt.strip(): return prompt suffix = random.choice(ENHANCE_SUFFIXES) base = prompt.rstrip(".").rstrip(",").strip() return base + suffix def on_lora_select(evt: gr.SelectData, aspect): lora = loras[evt.index] placeholder = f"Describe your image for {lora['title']}..." info_md = f"### [{lora['repo']}](https://huggingface.co/{lora['repo']}) ✅" new_aspect = aspect if "aspect" in lora: new_aspect = {"portrait": "9:16", "landscape": "16:9"}.get(lora["aspect"], aspect) return gr.update(placeholder=placeholder), info_md, evt.index, new_aspect def on_speed_change(speed): if speed == "Fast · 8 steps": return gr.update(value="Fast · 8 steps with Lightning LoRA"), 8, 1.0 return gr.update(value="Base · 50 steps — best quality"), 50, 4.0 def fetch_hf_lora(link: str): parts = link.strip("/").split("/") if len(parts) < 2: raise ValueError("Invalid repo path.") repo = "/".join(parts[-2:]) card = ModelCard.load(repo) base = card.data.get("base_model", "") bases = base if isinstance(base, list) else [base] if not any("Qwen/Qwen-Image" in b for b in bases): raise ValueError("Not a Qwen-Image LoRA.") trigger = card.data.get("instance_prompt", "") img_path = card.data.get("widget", [{}])[0].get("output", {}).get("url") image_url = f"https://huggingface.co/{repo}/resolve/main/{img_path}" if img_path else None fs = HfFileSystem() files = fs.ls(repo, detail=False) weight = next((f.split("/")[-1] for f in files if f.endswith(".safetensors")), None) if not weight: raise ValueError("No .safetensors found.") return parts[-1], repo, weight, trigger, image_url def add_custom_lora(custom_text: str): global loras if not custom_text.strip(): return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None try: link = custom_text.strip() if "huggingface.co/" in link: link = link.split("huggingface.co/")[-1] title, repo, weight, trigger, image = fetch_hf_lora(link) existing = next((i for i, l in enumerate(loras) if l["repo"] == repo), None) if existing is None: loras.append({"image": image, "title": title, "repo": repo, "weights": weight, "trigger_word": trigger}) existing = len(loras) - 1 card_html = f"""
{'' if image else ''}
{title}
{('trigger: '+trigger+'') if trigger else 'no trigger word'}
""" new_gallery = [(l["image"], l["title"]) for l in loras] return (gr.update(visible=True, value=card_html), gr.update(visible=True), gr.update(value=new_gallery, selected_index=None), f"Custom: {weight}", existing) except Exception as e: gr.Warning(str(e)) return gr.update(visible=True, value=f"{e}"), gr.update(visible=True), gr.update(), "", None def remove_custom_lora(): gallery_reset = [(l["image"], l["title"]) for l in loras] return gr.update(visible=False), gr.update(visible=False), gr.update(value=gallery_reset), "", None generate.zerogpu = True # ── CSS ─────────────────────────────────────────────────────────────────────── CSS = """ @import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap'); body, .gradio-container { background: #0b0a09 !important; font-family: 'Space Grotesk', sans-serif !important; } /* ── HEADER ── */ #drex-header { text-align: center; padding: 20px 0 8px; border-bottom: 1px solid #272422; margin-bottom: 16px; } #drex-header .drex-logo { display: inline-flex; align-items: center; gap: 12px; } #drex-header .drex-mark { width: 38px; height: 38px; background: #c2410c; border-radius: 8px; display: flex; align-items: center; justify-content: center; } #drex-header .drex-mark svg { width: 20px; height: 20px; fill: white; } #drex-header h1 { font-family: 'Space Grotesk', sans-serif !important; font-size: 28px !important; font-weight: 700 !important; color: #f0ebe5 !important; letter-spacing: -1px; margin: 0; } #drex-header .drex-sub { font-family: 'JetBrains Mono', monospace; font-size: 11px; color: #524e4a; margin-top: 3px; } /* ── THEME BAR ── */ #theme-bar { display: flex; align-items: center; justify-content: center; gap: 8px; padding: 10px 0 14px; } .theme-dot { width: 20px; height: 20px; border-radius: 50%; cursor: pointer; border: 2px solid transparent; transition: transform .15s; display: inline-block; } .theme-dot:hover { transform: scale(1.2); } /* ── BLOCKS ── */ .gradio-container .block { background: #131110 !important; border: 1px solid #272422 !important; border-radius: 10px !important; } label, .label-wrap span { font-family: 'JetBrains Mono', monospace !important; font-size: 10px !important; color: #524e4a !important; letter-spacing: 1px; text-transform: uppercase; } /* ── GALLERY ── */ #lora-gallery .grid-wrap { height: 200px !important; } #lora-gallery .thumbnail-item { border-radius: 6px !important; border: 1px solid #272422 !important; overflow: hidden; } #lora-gallery .thumbnail-item.selected { border-color: #c2410c !important; box-shadow: 0 0 0 1px #c2410c !important; } /* ── INPUTS ── */ textarea, input[type=text], input[type=number] { background: #1a1816 !important; border: 1px solid #272422 !important; border-radius: 6px !important; color: #f0ebe5 !important; font-family: 'Space Grotesk', sans-serif !important; } textarea:focus, input:focus { border-color: #c2410c !important; } /* ── BUTTONS ── */ button.primary { background: #c2410c !important; border: none !important; border-radius: 6px !important; font-family: 'Space Grotesk', sans-serif !important; font-weight: 600 !important; } button.primary:hover { filter: brightness(1.1); } button.secondary { background: #1a1816 !important; border: 1px solid #272422 !important; border-radius: 6px !important; color: #9a9088 !important; } button.secondary:hover { border-color: #c2410c !important; color: #c2410c !important; } /* ── SLIDERS ── */ input[type=range] { accent-color: #c2410c; } /* ── TABS ── */ .tab-nav button { font-family: 'JetBrains Mono', monospace !important; font-size: 11px !important; color: #524e4a !important; background: transparent !important; border: none !important; border-bottom: 2px solid transparent !important; border-radius: 0 !important; } .tab-nav button.selected { color: #c2410c !important; border-bottom-color: #c2410c !important; } /* ── STATUS ── */ #status-bar { font-family: 'JetBrains Mono', monospace !important; font-size: 11px !important; background: #131110 !important; border: 1px solid #272422 !important; border-radius: 6px !important; padding: 7px 12px !important; color: #9a9088 !important; } /* ── METADATA ── */ #meta-panel { font-family: 'JetBrains Mono', monospace !important; } /* ── GEN BTN ── */ #gen-btn { height: 44px !important; font-size: 14px !important; letter-spacing: 0.3px; } /* ── ACCORDION ── */ .accordion { background: #131110 !important; border: 1px solid #272422 !important; } .accordion .label-wrap { color: #9a9088 !important; } """ # ── JAVASCRIPT (theme switcher + prompt enhancer hint) ──────────────────────── JS_INIT = """ function() { // expose accent-color CSS var for metadata panel document.documentElement.style.setProperty('--drex-acc', '#c2410c'); } """ # ── GRADIO UI ───────────────────────────────────────────────────────────────── with gr.Blocks(title="D-REX Studio") as app: selected_index = gr.State(value=None) history_state = gr.State(value=[]) # ── HEADER ── gr.HTML("""
""") # ── LOGIN ── with gr.Row(): gr.LoginButton(scale=0) # ── THEME BAR ── with gr.Row(): theme_selector = gr.HTML("""
THEME CUSTOM
""") # ── MAIN ── with gr.Row(): # ── LEFT PANEL ── with gr.Column(scale=4): with gr.Tabs(): with gr.Tab("Gallery"): gallery = gr.Gallery( value=[(l["image"], l["title"]) for l in loras], label=None, allow_preview=False, columns=3, elem_id="lora-gallery", show_label=False, ) selected_info = gr.Markdown("", elem_id="lora-info") with gr.Tab("Favorites"): gr.Markdown("*Star a LoRA in the gallery to save it here.*", elem_id="fav-placeholder") fav_html = gr.HTML( "
no favorites yet
" ) gr.Markdown( "> Tip: Right-click a LoRA image → **Add to favorites** coming in next update.", elem_id="fav-tip" ) with gr.Tab("History"): history_html = gr.HTML( "
no history yet
", elem_id="history-display", ) clear_history_btn = gr.Button("Clear history", size="sm", variant="secondary") with gr.Tab("Custom LoRA"): custom_lora_input = gr.Textbox( label="HuggingFace repo", placeholder="username/lora-model-name", show_label=True, ) gr.Markdown("[Browse Qwen-Image LoRAs →](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image)") custom_lora_info = gr.HTML(visible=False) custom_lora_remove = gr.Button("Remove custom LoRA", visible=False, size="sm", variant="secondary") # ── RIGHT PANEL ── with gr.Column(scale=5): with gr.Row(): prompt = gr.Textbox( label="Prompt", placeholder="Describe your image...", lines=2, scale=5, ) with gr.Column(scale=1, min_width=90): enhance_btn = gr.Button("✦ Enhance", variant="secondary", size="sm") gen_btn = gr.Button("Generate", variant="primary", elem_id="gen-btn") neg_prompt = gr.Textbox( label="Negative prompt", placeholder="blur, watermark, low quality...", lines=1, ) result = gr.Image(label="Output", format="png", elem_id="output-image") meta_panel = gr.HTML( "
", elem_id="meta-panel", visible=True, ) 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="3:2", ) speed_mode = gr.Dropdown( label="Mode", choices=["Base · 50 steps", "Fast · 8 steps"], value="Base · 50 steps", ) status_bar = gr.Textbox( value="D-REX ready · base · 50 steps", label=None, interactive=False, show_label=False, elem_id="status-bar", ) with gr.Accordion("Advanced settings", open=False): with gr.Row(): cfg_scale = gr.Slider( label="CFG Scale", minimum=1.0, maximum=5.0, step=0.1, value=4.0 ) steps = gr.Slider( label="Steps", minimum=4, maximum=50, step=1, value=50 ) 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) # ── EVENTS ─────────────────────────────────────────────────────────────── gallery.select( on_lora_select, inputs=[aspect_ratio], outputs=[prompt, selected_info, selected_index, aspect_ratio], ) speed_mode.change( on_speed_change, inputs=[speed_mode], outputs=[status_bar, steps, cfg_scale], ) enhance_btn.click( enhance_prompt, inputs=[prompt], outputs=[prompt], ) gr.on( triggers=[gen_btn.click, prompt.submit], fn=generate, inputs=[ prompt, neg_prompt, cfg_scale, steps, selected_index, randomize_seed, seed, aspect_ratio, lora_scale, speed_mode, history_state, ], outputs=[result, seed, meta_panel, history_html, history_state], ) # Note: oauth_token injected automatically by Gradio from LoginButton custom_lora_input.input( add_custom_lora, inputs=[custom_lora_input], outputs=[custom_lora_info, custom_lora_remove, gallery, selected_info, selected_index], ) custom_lora_remove.click( remove_custom_lora, outputs=[custom_lora_info, custom_lora_remove, gallery, selected_info, selected_index], ) clear_history_btn.click( lambda: ([], "
history cleared
"), outputs=[history_state, history_html], ) # ── LAUNCH ──────────────────────────────────────────────────────────────────── app.queue() app.launch( theme=drex_theme, css=CSS, js=JS_INIT, )