Update app.py

app.py

@@ -1,329 +1,309 @@
 import os
-import random
-import uuid
-import json
-import time
-import asyncio
-import re
-from threading import Thread
-
 import gradio as gr
-import spaces
+import json
+import logging
 import torch
-import numpy as np
 from PIL import Image
-import cv2
-
-from transformers import (
-    AutoProcessor,
-    Gemma3ForConditionalGeneration,
-    Qwen2VLForConditionalGeneration,
-    TextIteratorStreamer,
-)
-from transformers.image_utils import load_image
+import random
+import time
+from hi_diffusers import HiDreamImagePipeline, HiDreamImageTransformer2DModel
+from hi_diffusers.schedulers.flash_flow_match import FlashFlowMatchEulerDiscreteScheduler
+from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
+from huggingface_hub import ModelCard
 
 # Constants
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
-MAX_SEED = np.iinfo(np.int32).max
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+MODEL_PREFIX = "HiDream-ai"
+LLAMA_MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 
-# Helper function to return a progress bar HTML snippet.
-def progress_bar_html(label: str) -> str:
-    return f'''
-<div style="display: flex; align-items: center;">
-    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
-    <div style="width: 110px; height: 5px; background-color: #F0FFF0; border-radius: 2px; overflow: hidden;">
-        <div style="width: 100%; height: 100%; background-color: #00FF00; animation: loading 1.5s linear infinite;"></div>
-    </div>
-</div>
-<style>
-@keyframes loading {{
-    0% {{ transform: translateX(-100%); }}
-    100% {{ transform: translateX(100%); }}
-}}
-</style>
-    '''
+FAST_MODEL_CONFIG = {
+        "path": f"{MODEL_PREFIX}/HiDream-I1-Full",
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "shift": 3.0,
+        "scheduler": FlowUniPCMultistepScheduler
+}
 
-# Qwen2-VL (for optional image inference)
+RESOLUTION_OPTIONS = [
+    "1024 × 1024 (Square)",
+    "768 × 1360 (Portrait)",
+    "1360 × 768 (Landscape)",
+    "880 × 1168 (Portrait)",
+    "1168 × 880 (Landscape)",
+    "1248 × 832 (Landscape)",
+    "832 × 1248 (Portrait)"
+]
 
-MODEL_ID_VL = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct" 
-processor = AutoProcessor.from_pretrained(MODEL_ID_VL, trust_remote_code=True)
-model_m = Qwen2VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_VL,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
-).to("cuda").eval()
+# Load LoRAs from JSON file (assumed to be compatible with Hi-Dream)
+with open('loras.json', 'r') as f:
+    loras = json.load(f)
 
-def clean_chat_history(chat_history):
-    cleaned = []
-    for msg in chat_history:
-        if isinstance(msg, dict) and isinstance(msg.get("content"), str):
-            cleaned.append(msg)
-    return cleaned
+device = "cuda" if torch.cuda.is_available() else "cpu"
+MAX_SEED = 2**32 - 1
 
-bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
-bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
-default_negative = os.getenv("default_negative", "")
+# Parse resolution string to height and width
+def parse_resolution(res_str):
+    mapping = {
+        "1024 × 1024": (1024, 1024),
+        "768 × 1360": (768, 1360),
+        "1360 × 768": (1360, 768),
+        "880 × 1168": (880, 1168),
+        "1168 × 880": (1168, 880),
+        "1248 × 832": (1248, 832),
+        "832 × 1248": (832, 1248)
+    }
+    for key, (h, w) in mapping.items():
+        if key in res_str:
+            return h, w
+    return 1024, 1024  # fallback
 
-def check_text(prompt, negative=""):
-    for i in bad_words:
-        if i in prompt:
-            return True
-    for i in bad_words_negative:
-        if i in negative:
-            return True
-    return False
+# Load the Hi-Dream Fast Model pipeline
+pipe, MODEL_CONFIG = None, None
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
+def load_fast_model():
+    global pipe, MODEL_CONFIG
+    config = FAST_MODEL_CONFIG
+    scheduler = config["scheduler"](
+        num_train_timesteps=1000,
+        shift=config["shift"],
+        use_dynamic_shifting=False
+    )
 
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(
+        LLAMA_MODEL_NAME,
+        use_fast=False
+    )
+    text_encoder = LlamaForCausalLM.from_pretrained(
+        LLAMA_MODEL_NAME,
+        output_hidden_states=True,
+        output_attentions=True,
+        torch_dtype=torch.bfloat16
+    ).to(device)
 
-dtype = torch.float16 if device.type == "cuda" else torch.float32
+    transformer = HiDreamImageTransformer2DModel.from_pretrained(
+        config["path"],
+        subfolder="transformer",
+        torch_dtype=torch.bfloat16
+    ).to(device)
 
+    pipe = HiDreamImagePipeline.from_pretrained(
+        config["path"],
+        scheduler=scheduler,
+        tokenizer_4=tokenizer,
+        text_encoder_4=text_encoder,
+        torch_dtype=torch.bfloat16
+    ).to(device, torch.bfloat16)
 
-# Gemma3 Model (default for text, image, & video inference)
+    pipe.transformer = transformer
+    MODEL_CONFIG = config
+    return pipe, config
 
-gemma3_model_id = "google/gemma-3-4b-it" #[or] Duplicate the space to use 12b
-gemma3_model = Gemma3ForConditionalGeneration.from_pretrained(
-    gemma3_model_id, device_map="auto"
-).eval()
-gemma3_processor = AutoProcessor.from_pretrained(gemma3_model_id)
+# Generate image
+def generate_image(prompt, resolution, seed, guidance_scale, num_inference_steps):
+    global pipe, MODEL_CONFIG
+    if pipe is None:
+        pipe, MODEL_CONFIG = load_fast_model()
 
-# VIDEO PROCESSING HELPER
+    height, width = parse_resolution(resolution)
+    if seed == -1 or seed is None:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(int(seed))
 
-def downsample_video(video_path):
-    vidcap = cv2.VideoCapture(video_path)
-    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
-    fps = vidcap.get(cv2.CAP_PROP_FPS)
-    frames = []
-    # Sample 10 evenly spaced frames.
-    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
-    for i in frame_indices:
-        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
-        success, image = vidcap.read()
-        if success:
-            # Convert from BGR to RGB and then to PIL Image.
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            pil_image = Image.fromarray(image)
-            timestamp = round(i / fps, 2)
-            frames.append((pil_image, timestamp))
-    vidcap.release()
-    return frames
+    result = pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        num_images_per_prompt=1,
+        generator=generator
+    )
 
-# MAIN GENERATION FUNCTION
+    return result.images[0], seed
 
-@spaces.GPU
-def generate(
-    input_dict: dict,
-    chat_history: list[dict],
-    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
-):
-    text = input_dict["text"]
-    files = input_dict.get("files", [])
-    lower_text = text.lower().strip()
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_name
 
-    # ----- Qwen2-VL branch (triggered with @qwen2-vl) -----
-    if lower_text.startswith("@qwen2-vl"):
-        prompt_clean = re.sub(r"@qwen2-vl", "", text, flags=re.IGNORECASE).strip().strip('"')
-        if files:
-            images = [load_image(f) for f in files]
-            messages = [{
-                "role": "user",
-                "content": [
-                    *[{"type": "image", "image": image} for image in images],
-                    {"type": "text", "text": prompt_clean},
-                ]
-            }]
-            prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-            inputs = processor(text=[prompt], images=images, return_tensors="pt", padding=True).to("cuda")
+    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:
-            messages = [
-                {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
-                {"role": "user", "content": [{"type": "text", "text": prompt_clean}]}
-            ]
-            inputs = processor.apply_chat_template(
-                messages, add_generation_prompt=True, tokenize=True,
-                return_dict=True, return_tensors="pt"
-            ).to("cuda", dtype=torch.float16)
-        streamer = TextIteratorStreamer(processor.tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = {
-            **inputs,
-            "streamer": streamer,
-            "max_new_tokens": max_new_tokens,
-            "do_sample": True,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repetition_penalty": repetition_penalty,
-        }
-        thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
-        thread.start()
-        buffer = ""
-        yield progress_bar_html("Processing with Qwen2VL")
-        for new_text in streamer:
-            buffer += new_text
-            buffer = buffer.replace("<|im_end|>", "")
-            time.sleep(0.01)
-            yield buffer
-        return
+            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
 
-    # ----- Default branch: Gemma3 (for text, image, & video inference) -----
-    if files:
-        # Check if any provided file is a video based on extension.
-        video_extensions = (".mp4", ".mov", ".avi", ".mkv", ".webm")
-        if any(str(f).lower().endswith(video_extensions) for f in files):
-            # Video inference branch.
-            prompt_clean = re.sub(r"@video-infer", "", text, flags=re.IGNORECASE).strip().strip('"')
-            video_path = files[0]
-            frames = downsample_video(video_path)
-            messages = [
-                {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
-                {"role": "user", "content": [{"type": "text", "text": prompt_clean}]}
-            ]
-            # Append each frame (with its timestamp) to the conversation.
-            for frame in frames:
-                image, timestamp = frame
-                image_path = f"video_frame_{uuid.uuid4().hex}.png"
-                image.save(image_path)
-                messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
-                messages[1]["content"].append({"type": "image", "url": image_path})
-            inputs = gemma3_processor.apply_chat_template(
-                messages, add_generation_prompt=True, tokenize=True,
-                return_dict=True, return_tensors="pt"
-            ).to(gemma3_model.device, dtype=torch.bfloat16)
-            streamer = TextIteratorStreamer(gemma3_processor.tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-            generation_kwargs = {
-                **inputs,
-                "streamer": streamer,
-                "max_new_tokens": max_new_tokens,
-                "do_sample": True,
-                "temperature": temperature,
-                "top_p": top_p,
-                "top_k": top_k,
-                "repetition_penalty": repetition_penalty,
-            }
-            thread = Thread(target=gemma3_model.generate, kwargs=generation_kwargs)
-            thread.start()
-            buffer = ""
-            yield progress_bar_html("Processing video with Gemma3")
-            for new_text in streamer:
-                buffer += new_text
-                time.sleep(0.01)
-                yield buffer
-            return
+def update_selection(evt: gr.SelectData, resolution):
+    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}) ✨"
+    if "aspect" in selected_lora:
+        if selected_lora["aspect"] == "portrait":
+            resolution = "768 × 1360 (Portrait)"
+        elif selected_lora["aspect"] == "landscape":
+            resolution = "1360 × 768 (Landscape)"
         else:
-            # Image inference branch.
-            prompt_clean = re.sub(r"@gemma3", "", text, flags=re.IGNORECASE).strip().strip('"')
-            images = [load_image(f) for f in files]
-            messages = [{
-                "role": "user",
-                "content": [
-                    *[{"type": "image", "image": image} for image in images],
-                    {"type": "text", "text": prompt_clean},
-                ]
-            }]
-            inputs = gemma3_processor.apply_chat_template(
-                messages, tokenize=True, add_generation_prompt=True,
-                return_dict=True, return_tensors="pt"
-            ).to(gemma3_model.device, dtype=torch.bfloat16)
-            streamer = TextIteratorStreamer(gemma3_processor.tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-            generation_kwargs = {
-                **inputs,
-                "streamer": streamer,
-                "max_new_tokens": max_new_tokens,
-                "do_sample": True,
-                "temperature": temperature,
-                "top_p": top_p,
-                "top_k": top_k,
-                "repetition_penalty": repetition_penalty,
-            }
-            thread = Thread(target=gemma3_model.generate, kwargs=generation_kwargs)
-            thread.start()
-            buffer = ""
-            yield progress_bar_html("Processing with Gemma3")
-            for new_text in streamer:
-                buffer += new_text
-                time.sleep(0.01)
-                yield buffer
-            return
-    else:
-        # Text-only inference branch.
-        messages = [
-            {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
-            {"role": "user", "content": [{"type": "text", "text": text}]}
-        ]
-        inputs = gemma3_processor.apply_chat_template(
-            messages, add_generation_prompt=True, tokenize=True,
-            return_dict=True, return_tensors="pt"
-        ).to(gemma3_model.device, dtype=torch.bfloat16)
-        streamer = TextIteratorStreamer(gemma3_processor.tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = {
-            **inputs,
-            "streamer": streamer,
-            "max_new_tokens": max_new_tokens,
-            "do_sample": True,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repetition_penalty": repetition_penalty,
-        }
-        thread = Thread(target=gemma3_model.generate, kwargs=generation_kwargs)
-        thread.start()
-        outputs = []
-        for new_text in streamer:
-            outputs.append(new_text)
-            yield "".join(outputs)
-        final_response = "".join(outputs)
-        yield final_response
+            resolution = "1024 × 1024 (Square)"
+    return (
+        gr.update(placeholder=new_placeholder),
+        updated_text,
+        evt.index,
+        resolution,
+    )
+
+def run_lora(prompt, resolution, cfg_scale, steps, selected_index, randomize_seed, seed):
+    global pipe
+    if pipe is None:
+        pipe, _ = load_fast_model()
+
+    if selected_index is not None:
+        selected_lora = loras[selected_index]
+        lora_path = selected_lora["repo"]
+        weight_name = selected_lora.get("weights", None)
+        with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
+            pipe.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True)
+        trigger_word = selected_lora.get("trigger_word", "")
+        if trigger_word:
+            if "trigger_position" in selected_lora and selected_lora["trigger_position"] == "prepend":
+                prompt = f"{trigger_word} {prompt}"
+            else:
+                prompt = f"{prompt} {trigger_word}"
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    with calculateDuration("Generating image"):
+        final_image, used_seed = generate_image(prompt, resolution, seed, cfg_scale, steps)
+    return final_image, used_seed
+
+def check_custom_model(link):
+    split_link = link.split("/")
+    if len(split_link) != 2:
+        raise Exception("Invalid Hugging Face repository link format.")
+    model_card = ModelCard.load(link)
+    image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+    trigger_word = model_card.data.get("instance_prompt", "")
+    image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+    safetensors_name = None  # Simplified; assumes a safetensors file exists
+    return split_link[1], link, safetensors_name, trigger_word, image_url
+
+def add_custom_lora(custom_lora):
+    global loras
+    if custom_lora:
+        try:
+            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
+            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."}</small>
+                </div>
+              </div>
+            </div>
+            '''
+            existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
+            if not existing_item_index:
+                new_item = {
+                    "image": image,
+                    "title": title,
+                    "repo": repo,
+                    "weights": path,
+                    "trigger_word": trigger_word
+                }
+                existing_item_index = len(loras)
+                loras.append(new_item)
+            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: {str(e)}")
+            return gr.update(visible=True, value=f"Invalid LoRA: {str(e)}"), gr.update(visible=True), gr.update(), "", None, ""
+    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
+
+def remove_custom_lora():
+    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}
+'''
 
+font = [gr.themes.GoogleFont("Source Sans Pro"), "Arial", "sans-serif"]
+with gr.Blocks(theme=gr.themes.Soft(font=font), css=css, delete_cache=(60, 60)) as app:
+    title = gr.HTML(
+        """<h1>Hi-Dream Full 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="Type a prompt after selecting a LoRA")
+        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.Group():
+                custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path", placeholder="multimodalart/vintage-ads-flux")
+                gr.Markdown("[Check the list of Hi-Dream LoRAs]", 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")
 
-# Gradio Interface
+    with gr.Row():
+        with gr.Accordion("Advanced Settings", open=False):
+            cfg_scale = gr.Slider(label="Guidance Scale", minimum=0, maximum=20, step=0.1, value=FAST_MODEL_CONFIG["guidance_scale"])
+            steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=FAST_MODEL_CONFIG["num_inference_steps"])
+            resolution = gr.Radio(
+                choices=RESOLUTION_OPTIONS,
+                value=RESOLUTION_OPTIONS[0],
+                label="Resolution"
+            )
+            randomize_seed = gr.Checkbox(True, label="Randomize seed")
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
 
-demo = gr.ChatInterface(
-    fn=generate,
-    additional_inputs=[
-        gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
-        gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
-        gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50),
-        gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2),
-    ],
-    examples=[
-        [{"text": "Create a short story based on the image.","files": ["examples/1111.jpg"]}],
-        [{"text": "Explain the Image", "files": ["examples/3.jpg"]}],
-        [{"text": "Explain the content of the Advertisement", "files": ["examples/videoplayback.mp4"]}],
-        [{"text": "Which movie character is this?", "files": ["examples/9999.jpg"]}],
-        ["Explain Critical Temperature of Substance"],
-        [{"text": "@qwen2-vl Transcription of the letter", "files": ["examples/222.png"]}],
-        [{"text": "Explain the content of the video in detail", "files": ["examples/breakfast.mp4"]}],
-        [{"text": "Describe the video", "files": ["examples/Missing.mp4"]}],
-        [{"text": "Explain what is happening in this video ?", "files": ["examples/oreo.mp4"]}],
-        [{"text": "Summarize the events in this video", "files": ["examples/sky.mp4"]}],
-        [{"text": "What is in the video ?", "files": ["examples/redlight.mp4"]}],
-        ["Python Program for Array Rotation"],
-        ["Explain Critical Temperature of Substance"]
-    ],
-    cache_examples=False,
-    type="messages",
-    description="# **Gemma 3 Multimodal** \n`Use @qwen2-vl to switch to Qwen2-VL OCR for image inference and @video-infer for video input`",
-    fill_height=True,
-    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image", "video"], file_count="multiple", placeholder="Tag with @qwen2-vl for Qwen2-VL inference if needed."),
-    stop_btn="Stop Generation",
-    multimodal=True,
-)
+    gallery.select(
+        update_selection,
+        inputs=[resolution],
+        outputs=[prompt, selected_info, selected_index, resolution]
+    )
+    custom_lora.input(
+        add_custom_lora,
+        inputs=[custom_lora],
+        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
+    )
+    custom_lora_button.click(
+        remove_custom_lora,
+        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
+    )
+    gr.on(
+        triggers=[generate_button.click, prompt.submit],
+        fn=run_lora,
+        inputs=[prompt, resolution, cfg_scale, steps, selected_index, randomize_seed, seed],
+        outputs=[result, seed]
+    )
 
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch(share=True)
+app.queue()
+app.launch()