Upload app.py

app.py

@@ -0,0 +1,212 @@
+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()