Update app.py

app.py

@@ -1,212 +1,209 @@
-import spaces
-import os
-import gradio as gr
-import torch
-import numpy as np
-import random
-import requests
-import re
-from diffusers import FluxPipeline
-from translatepy import Translator
-from huggingface_hub import hf_hub_download
-
-# Environment setup
-os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-translator = Translator()
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-
-# Constants
-MODEL_ID = "black-forest-labs/FLUX.1-dev"
-DEFAULT_LORA = "nftnik/BR_ohwx_V1"
-DEFAULT_WEIGHT_NAME = "BR_ohwx.safetensors"
-MAX_SEED = np.iinfo(np.int32).max
-
-CSS = """
-footer {
-    visibility: hidden;
-}
-"""
-
-JS = """function () {
-  gradioURL = window.location.href;
-  if (!gradioURL.endsWith('?__theme=dark')) {
-    window.location.replace(gradioURL + '?__theme=dark');
-  }
-}"""
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using {device.upper()}")
-
-# Initialize pipeline and load default LoRA weights
-pipe = FluxPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to(device)
-pipe.load_lora_weights(DEFAULT_LORA, weight_name=DEFAULT_WEIGHT_NAME)
-
-def scrape_lora_link(url: str):
-    try:
-        response = requests.get(url)
-        response.raise_for_status()
-        content = response.text
-        pattern = r'href="(.*?lora.*?\.safetensors\?download=true)"'
-        pattern2 = r'href="(.*?\.safetensors\?download=true)"'
-        match = re.search(pattern, content)
-        match2 = re.search(pattern2, content)
-        if match:
-            safetensors_url = match.group(1)
-            filename = safetensors_url.split('/')[-1].split('?')[0]
-            return filename
-        elif match2:
-            safetensors_url = match2.group(1)
-            filename = safetensors_url.split('/')[-1].split('?')[0]
-            return filename
-        else:
-            return None
-    except requests.RequestException as e:
-        raise gr.Error(f"An error occurred while fetching the URL: {e}")
-
-def enable_lora(lora_add: str, progress=gr.Progress(track_tqdm=True)):
-    pipe.unload_lora_weights()
-    if not lora_add:
-        gr.Info("No LoRA Loaded, using base model")
-        return gr.update(value="")
-    else:
-        url = f"https://huggingface.co/{lora_add}/tree/main"
-        lora_name = scrape_lora_link(url)
-        if lora_name:
-            print(f"Loading LoRA: {lora_add}/{lora_name}")
-            pipe.load_lora_weights(lora_add, weight_name=lora_name)
-            gr.Info(f"{lora_add} Loaded")
-            return gr.update(label="LoRA Loaded Now")
-        else:
-            try:
-                pipe.load_lora_weights(lora_add)
-                print(f"Loading LoRA: {lora_add}")
-                gr.Info(f"{lora_add} Loaded")
-                return gr.update(label="LoRA Loaded Now")
-            except Exception as e:
-                raise gr.Error(f"{lora_add} load failed: {e}")
-
-# Placeholder function to update flux scheduler and sampler settings.
-def update_flux_settings(scheduler_choice: str, sampler_choice: str):
-    # Replace the code below with actual logic to update your pipeline's scheduler/sampler.
-    print(f"Setting scheduler to {scheduler_choice} and sampler to {sampler_choice}")
-    # e.g.:
-    # pipe.scheduler = SchedulerClassMapping[scheduler_choice].from_config(pipe.scheduler.config)
-    # pipe.sampler = SamplerClassMapping[sampler_choice](**pipe.sampler_config)
-    return f"Scheduler set to {scheduler_choice} and Sampler set to {sampler_choice}"
-
-@spaces.GPU()
-def generate_image(
-    prompt: str,
-    lora_word: str,
-    lora_scale: float = 0.8,
-    width: int = 896,
-    height: int = 1152,
-    guidance_scale: float = 3.5,
-    steps: int = 25,
-    seed: int = -1,
-    nums: int = 1,
-    progress=gr.Progress(track_tqdm=True)
-):
-    # Ensure the pipeline is on the correct device.
-    pipe.to(device)
-    if seed == -1:
-        seed = random.randint(0, MAX_SEED)
-    seed = int(seed)
-    
-    # Translate prompt to English.
-    translation = translator.translate(prompt, "English")
-    prompt_english = str(translation)  # Adjust if translatepy returns a different attribute.
-    full_prompt = f"{prompt_english} {lora_word}"
-    print(f"Prompt: {full_prompt}")
-    
-    generator = torch.Generator().manual_seed(seed)
-    result = pipe(
-        prompt=full_prompt,
-        height=height,
-        width=width,
-        guidance_scale=guidance_scale,
-        output_type="pil",
-        num_inference_steps=steps,
-        max_sequence_length=512,
-        num_images_per_prompt=nums,
-        generator=generator,
-        joint_attention_kwargs={"scale": lora_scale},
-    )
-    return result.images, seed
-
-# Sample examples
-examples = [
-    ["close-up portrait of a futuristic alien in cyberpunk attire", "ohwx", 0.9],
-    ["full-body shot of an alien running through a neon-lit cityscape", "ohwx", 0.9],
-    ["portrait of a blue alien with techwear in a virtual reality environment", "ohwx", 0.9],
-    ["cyberpunk style portrait of an alien with glowing eyes", "ohwx", 0.9]
-]
-
-with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
-    gr.HTML("<h1><center>BR METAVERSO - Avatar Generator</center></h1>")
-    gr.HTML("<p><center>Load the LoRA model on the menu</center></p>")
-    with gr.Row():
-        with gr.Column(scale=4):
-            gallery = gr.Gallery(label="Flux Generated Image", columns=1, preview=True, height=600)
-            with gr.Row():
-                prompt_input = gr.Textbox(
-                    label="Enter Your Prompt (Multi-Languages)", 
-                    lines=2, 
-                    placeholder="Enter prompt...", 
-                    scale=6
-                )
-                generate_btn = gr.Button(scale=1, variant="primary")
-        with gr.Accordion("Advanced Options", open=True):
-            with gr.Column(scale=1):
-                width_slider = gr.Slider(label="Width", minimum=512, maximum=1920, step=8, value=896)
-                height_slider = gr.Slider(label="Height", minimum=512, maximum=1920, step=8, value=1152)
-                guidance_slider = gr.Slider(label="Guidance Scale", minimum=3.5, maximum=7, step=0.1, value=3.5)
-                steps_slider = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=25)
-                seed_slider = gr.Slider(label="Seed", minimum=-1, maximum=MAX_SEED, step=1, value=-1)
-                nums_slider = gr.Slider(label="Image Count", minimum=1, maximum=4, step=1, value=1)
-            with gr.Column(scale=1):
-                lora_scale_slider = gr.Slider(label="LoRA Scale", minimum=0.1, maximum=2.0, step=0.1, value=1.0)
-                lora_add_text = gr.Textbox(
-                    label="Flux LoRA", 
-                    info="Copy the HF LoRA model name here", 
-                    lines=1, 
-                    value="nftnik/BR_ohwx_V1"
-                )
-                lora_word_text = gr.Textbox(
-                    label="Flux LoRA Trigger Word", 
-                    info="Add the Trigger Word", 
-                    lines=1, 
-                    value="ohwx"
-                )
-                load_lora_btn = gr.Button(value="Load LoRA", variant="secondary")
-                # New dropdowns for flux scheduler and sampler
-                flux_scheduler = gr.Dropdown(
-                    label="Flux Scheduler", 
-                    choices=["DDIM", "PNDM", "DPMSolver"], 
-                    value="DDIM"
-                )
-                flux_sampler = gr.Dropdown(
-                    label="Flux Sampler", 
-                    choices=["Default", "k_euler", "k_lms"], 
-                    value="Default"
-                )
-                update_flux_btn = gr.Button(value="Update Flux Settings", variant="secondary")
-                flux_status = gr.Textbox(label="Flux Settings Status", interactive=False)
-                
-    gr.Examples(
-        examples=examples,
-        inputs=[prompt_input, lora_word_text, lora_scale_slider],
-        cache_examples=False,
-        examples_per_page=4,
-    )
-
-    load_lora_btn.click(fn=enable_lora, inputs=[lora_add_text], outputs=lora_add_text)
-    update_flux_btn.click(fn=update_flux_settings, inputs=[flux_scheduler, flux_sampler], outputs=flux_status)
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt_input, lora_word_text, lora_scale_slider, width_slider, height_slider, guidance_slider, steps_slider, seed_slider, nums_slider],
-        outputs=[gallery, seed_slider],
-        api_name="run",
-    )
-    
-demo.queue().launch()
+import spaces
+import os
+import gradio as gr
+import torch
+import numpy as np
+import random
+import requests
+import re
+from diffusers import FluxPipeline
+from translatepy import Translator
+from huggingface_hub import hf_hub_download
+
+# Environment setup
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+translator = Translator()
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+# Constants and configuration
+MODEL_ID = "black-forest-labs/FLUX.1-dev"
+DEFAULT_LORA = "nftnik/BR_ohwx_V1"
+DEFAULT_WEIGHT_NAME = "BR_ohwx.safetensors"
+MAX_SEED = np.iinfo(np.int32).max
+
+CSS = """
+footer {
+    visibility: hidden;
+}
+"""
+
+JS = """function () {
+  const gradioURL = window.location.href;
+  if (!gradioURL.endsWith('?__theme=dark')) {
+    window.location.replace(gradioURL + '?__theme=dark');
+  }
+}"""
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using {device.upper()}")
+
+# Initialize the Flux pipeline
+pipe = FluxPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to(device)
+
+# Set the default sampler and scheduler.
+# NOTE: This example assumes that your FluxPipeline has methods `set_sampler` and `set_scheduler`
+# which accept a string indicating the desired method.
+print("Setting default sampler to 'euler' and default scheduler to 'beta' ...")
+pipe.set_sampler("euler")   # Replace with the correct call if your API differs
+pipe.set_scheduler("beta")  # Replace with the correct call if your API differs
+
+# Load the default LoRA weights
+pipe.load_lora_weights(DEFAULT_LORA, weight_name=DEFAULT_WEIGHT_NAME)
+
+def scrape_lora_link(url: str):
+    try:
+        response = requests.get(url)
+        response.raise_for_status()
+        content = response.text
+        pattern = r'href="(.*?lora.*?\.safetensors\?download=true)"'
+        pattern2 = r'href="(.*?\.safetensors\?download=true)"'
+        match = re.search(pattern, content)
+        match2 = re.search(pattern2, content)
+        if match:
+            safetensors_url = match.group(1)
+            filename = safetensors_url.split('/')[-1].split('?')[0]
+            return filename
+        elif match2:
+            safetensors_url = match2.group(1)
+            filename = safetensors_url.split('/')[-1].split('?')[0]
+            return filename
+        else:
+            return None
+    except requests.RequestException as e:
+        raise gr.Error(f"An error occurred while fetching the URL: {e}")
+
+def enable_lora(lora_add: str, progress=gr.Progress(track_tqdm=True)):
+    pipe.unload_lora_weights()
+    if not lora_add:
+        gr.Info("No LoRA Loaded, using base model")
+        return gr.update(value="")
+    else:
+        url = f"https://huggingface.co/{lora_add}/tree/main"
+        lora_name = scrape_lora_link(url)
+        if lora_name:
+            print(f"Loading LoRA: {lora_add}/{lora_name}")
+            pipe.load_lora_weights(lora_add, weight_name=lora_name)
+            gr.Info(f"{lora_add} Loaded")
+            return gr.update(label="LoRA Loaded Now")
+        else:
+            try:
+                pipe.load_lora_weights(lora_add)
+                print(f"Loading LoRA: {lora_add}")
+                gr.Info(f"{lora_add} Loaded")
+                return gr.update(label="LoRA Loaded Now")
+            except Exception as e:
+                raise gr.Error(f"{lora_add} load failed: {e}")
+
+@spaces.GPU()
+def generate_image(
+    prompt: str,
+    lora_word: str,
+    lora_scale: float = 0.8,
+    width: int = 896,
+    height: int = 1152,
+    guidance_scale: float = 3.5,
+    steps: int = 25,
+    seed: int = -1,
+    nums: int = 1,
+    progress=gr.Progress(track_tqdm=True)
+):
+    # Make sure the pipeline is on the correct device.
+    pipe.to(device)
+    if seed == -1:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    
+    # Translate the prompt to English.
+    translation = translator.translate(prompt, "English")
+    prompt_english = str(translation)  # Adjust if translatepy returns a different attribute
+    full_prompt = f"{prompt_english} {lora_word}"
+    print(f"Prompt: {full_prompt}")
+    
+    generator = torch.Generator().manual_seed(seed)
+    result = pipe(
+        prompt=full_prompt,
+        height=height,
+        width=width,
+        guidance_scale=guidance_scale,
+        output_type="pil",
+        num_inference_steps=steps,
+        max_sequence_length=512,
+        num_images_per_prompt=nums,
+        generator=generator,
+        joint_attention_kwargs={"scale": lora_scale},
+    )
+    return result.images, seed
+
+# Example prompts for demonstration
+examples = [
+    ["Medium-shot portrait, ohwx blue alien, wearing black techwear with a high collar, standing inside a futuristic VR showroom.  The background consists of a modern, clean showroom with diffused color neon lighting, creating a high-end, sophisticated aesthetic", "ohwx", 0.9],
+    ["ohwx blue alien, wearing black techwear with a high collar, immersed in a digital cybernetic landscape. A neon purple and magenta glow reflects on his skin as he stands inside a Metaverse oasis. His expression is focused, with his hands outstretched, interacting with the ambient. The environment is sleek and futuristic, with deep shadows and vibrant lighting creating a cinematic composition", "ohwx", 0.9],
+    ["full-body shot, ohwx blue alien, wearing black techwear with a high collar, black cyber sneakers, running through a neon-lit cyberpunk alley at night. His posture is dynamic, mid-stride, arms pumping. The wet pavement reflects the bright neon signs from above, casting colorful reflections on his sleek techwear. The deep shadows and dramatic lighting emphasize the futuristic setting", "ohwx", 0.9],
+    ["ohwx blue alien, wearing black techwear with a high collar, sitting inside a sleek, high-tech VR capsule, immersed in an augmented reality experience. His eyes glow with digital overlays. The lighting is a mix of deep grey ambient hues with bright cyan highlights from the AR projections.", "ohwx", 0.9]
+]
+
+# Build the Gradio interface
+with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
+    gr.HTML("<h1><center>BR METAVERSO - Avatar Generator</center></h1>")
+    gr.HTML("<p><center>Load the LoRA model on the menu</center></p>")
+    with gr.Row():
+        with gr.Column(scale=4):
+            gallery = gr.Gallery(label="Flux Generated Image", columns=1, preview=True, height=600)
+            with gr.Row():
+                prompt_input = gr.Textbox(
+                    label="Enter Your Prompt (Multi-Languages)",
+                    lines=2,
+                    placeholder="Enter prompt...",
+                    scale=6
+                )
+                generate_btn = gr.Button(scale=1, variant="primary")
+        with gr.Accordion("Advanced Options", open=True):
+            with gr.Column(scale=1):
+                width_slider = gr.Slider(label="Width", minimum=512, maximum=1920, step=8, value=896)
+                height_slider = gr.Slider(label="Height", minimum=512, maximum=1920, step=8, value=1152)
+                guidance_slider = gr.Slider(label="Guidance Scale", minimum=3.5, maximum=7, step=0.1, value=3.5)
+                steps_slider = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=25)
+                seed_slider = gr.Slider(label="Seed", minimum=-1, maximum=MAX_SEED, step=1, value=-1)
+                nums_slider = gr.Slider(label="Image Count", minimum=1, maximum=4, step=1, value=1)
+            with gr.Column(scale=1):
+                lora_scale_slider = gr.Slider(label="LoRA Scale", minimum=0.1, maximum=2.0, step=0.1, value=1.0)
+                lora_add_text = gr.Textbox(
+                    label="Flux LoRA",
+                    info="Copy the HF LoRA model name here",
+                    lines=1,
+                    value="nftnik/BR_ohwx_V1"
+                )
+                lora_word_text = gr.Textbox(
+                    label="Flux LoRA Trigger Word",
+                    info="Add the Trigger Word",
+                    lines=1,
+                    value="ohwx"
+                )
+                load_lora_btn = gr.Button(value="Load LoRA", variant="secondary")
+
+    gr.Examples(
+        examples=examples,
+        inputs=[prompt_input, lora_word_text, lora_scale_slider],
+        cache_examples=False,
+        examples_per_page=4,
+    )
+
+    load_lora_btn.click(fn=enable_lora, inputs=[lora_add_text], outputs=lora_add_text)
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[
+            prompt_input,
+            lora_word_text,
+            lora_scale_slider,
+            width_slider,
+            height_slider,
+            guidance_slider,
+            steps_slider,
+            seed_slider,
+            nums_slider
+        ],
+        outputs=[gallery, seed_slider],
+        api_name="run",
+    )
+    
+demo.queue().launch(ssr=False)