Update app.py

app.py

@@ -1,290 +1,132 @@
 import os
 import random
+import re
+import requests
 import torch
 import numpy as np
 import gradio as gr
+import spaces
 from diffusers import FluxPipeline
 from translatepy import Translator
 
 # -----------------------------------------------------------------------------
 # CONFIGURATION
 # -----------------------------------------------------------------------------
-class Config:
-    MODEL_ID = "black-forest-labs/FLUX.1-dev"
-    DEFAULT_LORA = "nftnik/BR_ohwx_V1"
-    DEFAULT_WEIGHT_NAME = "BR_ohwx.safetensors"
-    MAX_SEED = int(np.iinfo(np.int32).max)
-    CSS = "footer { visibility: hidden; }"
-    DEFAULT_WIDTH = 896
-    DEFAULT_HEIGHT = 1152
-    DEFAULT_GUIDANCE_SCALE = 3.5
-    DEFAULT_STEPS = 35
-    DEFAULT_LORA_SCALE = 1.0
-    DEFAULT_TRIGGER_WORD = "ohwx"
-    # Memory optimization configs
-    ENABLE_MEMORY_EFFICIENT_ATTENTION = True
-    ENABLE_SEQUENTIAL_CPU_OFFLOAD = True
-    ENABLE_ATTENTION_SLICING = "max"
+config = {
+    "model_id": "black-forest-labs/FLUX.1-dev",
+    "default_lora": "nftnik/BR_ohwx_V1",
+    "default_weight_name": "BR_ohwx.safetensors",
+    "max_seed": int(np.iinfo(np.int32).max),
+    "css": "footer { visibility: hidden; }",
+    "default_width": 896,
+    "default_height": 1152,
+    "default_guidance_scale": 3.5,
+    "default_steps": 35,
+    "default_loRa_scale": 1.0,
+}
 
 # -----------------------------------------------------------------------------
-# FluxGenerator class to handle image generation
+# Environment and device setup
 # -----------------------------------------------------------------------------
-class FluxGenerator:
-    def __init__(self):
-        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-        self.translator = Translator()
-        self.device = self._get_optimal_device()
-        print(f"Using {self.device.upper()}")
-
-        self.pipe = None
-        self._initialize_pipeline()
-
-    def _get_optimal_device(self):
-        """Determine the optimal device based on available resources"""
-        if torch.cuda.is_available():
-            try:
-                gpu_memory = torch.cuda.get_device_properties(0).total_memory
-                if gpu_memory > 10 * 1024 * 1024 * 1024:  # More than 10GB
-                    return "cuda"
-                else:
-                    print("Limited GPU memory detected. Will still use CUDA with memory optimizations.")
-                    return "cuda"
-            except:
-                print("Error checking GPU memory, falling back to CPU")
-                return "cpu"
-        else:
-            return "cpu"
-
-    def _initialize_pipeline(self):
-        """Initialize the Flux pipeline with memory optimizations"""
-        try:
-            print("Loading Flux model...")
-            pipe_kwargs = {
-                "torch_dtype": torch.bfloat16 if self.device == "cuda" else torch.float32,
-            }
-            self.pipe = FluxPipeline.from_pretrained(Config.MODEL_ID, **pipe_kwargs)
-
-            # Apply memory optimizations
-            if Config.ENABLE_MEMORY_EFFICIENT_ATTENTION and self.device == "cuda":
-                print("Enabling memory efficient attention")
-                self.pipe.enable_xformers_memory_efficient_attention()
-
-            if Config.ENABLE_ATTENTION_SLICING:
-                print("Enabling attention slicing")
-                self.pipe.enable_attention_slicing(Config.ENABLE_ATTENTION_SLICING)
-
-            if Config.ENABLE_SEQUENTIAL_CPU_OFFLOAD and self.device == "cuda":
-                print("Enabling sequential CPU offload")
-                self.pipe.enable_sequential_cpu_offload()
-            else:
-                # Only move to device if not offloading
-                self.pipe.to(self.device)
-
-            print(f"Loading default LoRA: {Config.DEFAULT_LORA}")
-            self.pipe.load_lora_weights(Config.DEFAULT_LORA, weight_name=Config.DEFAULT_WEIGHT_NAME)
-
-            print("Model initialization complete")
-        except Exception as e:
-            error_msg = f"Error initializing pipeline: {str(e)}"
-            print(error_msg)
-            raise
-
-    def load_lora(self, lora_path):
-        """Load a new LoRA model"""
-        try:
-            print("Unloading previous LoRA weights...")
-            self.pipe.unload_lora_weights()
-
-            if not lora_path:
-                print("No LoRA path provided, skipping LoRA loading.")
-                return gr.update(value="")
-
-            print(f"Loading LoRA from {lora_path}...")
-            self.pipe.load_lora_weights(lora_path)
-            print("LoRA loaded successfully.")
-            return gr.update(label="LoRA Loaded Successfully")
-        except Exception as e:
-            error_msg = f"Failed to load LoRA from {lora_path}: {str(e)}"
-            print(error_msg)
-            raise gr.Error(error_msg)
-
-    def _clear_memory(self):
-        """Clear CUDA memory cache"""
-        if self.device == "cuda":
-            try:
-                print("Clearing CUDA memory cache...")
-                torch.cuda.empty_cache()
-                if hasattr(torch.cuda, "amp") and hasattr(torch.cuda.amp, "autocast"):
-                    torch.cuda.amp.clear_autocast_cache()
-            except Exception as e:
-                print(f"Warning: Failed to clear CUDA memory: {str(e)}")
-
-    def generate(self, prompt, lora_word, lora_scale=Config.DEFAULT_LORA_SCALE,
-                 width=Config.DEFAULT_WIDTH, height=Config.DEFAULT_HEIGHT,
-                 guidance_scale=Config.DEFAULT_GUIDANCE_SCALE,
-                 steps=Config.DEFAULT_STEPS, seed=-1, num_images=1):
-        """Generate images from a prompt with memory optimizations."""
-        try:
-            print(f"Generating image for prompt: '{prompt}'")
-            self._clear_memory()
-
-            if not Config.ENABLE_SEQUENTIAL_CPU_OFFLOAD:
-                print(f"Moving model to {self.device}")
-                self.pipe.to(self.device)
-
-            seed = random.randint(0, Config.MAX_SEED) if seed == -1 else int(seed)
-            print(f"Using seed: {seed}")
-            generator = torch.Generator(device=self.device).manual_seed(seed)
-
-            print("Translating prompt if needed...")
-            prompt_english = str(self.translator.translate(prompt, "English"))
-            full_prompt = f"{prompt_english} {lora_word}"
-            print(f"Full prompt: '{full_prompt}'")
-
-            # If GPU memory is less than 8GB, scale resolution
-            if (self.device == "cuda" and
-                torch.cuda.get_device_properties(0).total_memory < 8 * 1024 * 1024 * 1024):
-                original_width, original_height = width, height
-                width = int(width * 0.85)
-                height = int(height * 0.85)
-                print(f"Memory is tight. Scaled resolution from {original_width}x{original_height} to {width}x{height}")
-
-            print(f"Starting generation with {steps} steps, guidance scale {guidance_scale}")
-            with torch.autocast("cuda", enabled=(self.device == "cuda")):
-                result = self.pipe(
-                    prompt=full_prompt,
-                    height=height,
-                    width=width,
-                    guidance_scale=guidance_scale,
-                    output_type="pil",
-                    num_inference_steps=steps,
-                    num_images_per_prompt=num_images,
-                    generator=generator,
-                    joint_attention_kwargs={"scale": lora_scale},
-                )
-            print("Generation complete, returning images.")
-            self._clear_memory()
-            return result.images, seed
-        except Exception as e:
-            error_msg = f"Image generation failed: {str(e)}"
-            print(error_msg)
-            self._clear_memory()
-            raise gr.Error(error_msg)
-
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+translator = Translator()
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using {device.upper()}")
 
 # -----------------------------------------------------------------------------
-# UI Builder class
+# Initialize the Flux pipeline and load default LoRA
 # -----------------------------------------------------------------------------
-class FluxUI:
-    def __init__(self, generator):
-        self.generator = generator
-        self.example_prompts = [
-            ["Medium-shot portrait, ohwx blue alien, wearing black techwear with a high collar, standing inside a futuristic VR showroom.", "ohwx", 0.9],
-            ["ohwx blue alien, wearing black techwear with a high collar, immersed in a digital cybernetic landscape.", "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.", "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.", "ohwx", 0.9]
-        ]
-
-    def build(self):
-        with gr.Blocks(css=Config.CSS) as demo:
-            gr.HTML("<h1><center>BR METAVERSO - Avatar Generator</center></h1>")
-
-            status_markdown = gr.Markdown("**🟢 Ready**", visible=True)
-
-            with gr.Row():
-                with gr.Column(scale=4):
-                    gallery = gr.Gallery(label="Flux Generated Image", columns=1, preview=True, height=600)
-                    prompt_input = gr.Textbox(
-                        label="Enter Your Prompt",
-                        lines=2,
-                        placeholder="Type your avatar description..."
-                    )
-                    generate_btn = gr.Button(value="Generate", variant="primary")
-
-                with gr.Accordion("Advanced Options", open=True):
-                    with gr.Row():
-                        with gr.Column():
-                            width_slider = gr.Slider(label="Width", minimum=512, maximum=1920, step=8, value=Config.DEFAULT_WIDTH)
-                            height_slider = gr.Slider(label="Height", minimum=512, maximum=1920, step=8, value=Config.DEFAULT_HEIGHT)
-                        with gr.Column():
-                            guidance_slider = gr.Slider(label="Guidance Scale", minimum=3.5, maximum=7, step=0.1, value=Config.DEFAULT_GUIDANCE_SCALE)
-                            steps_slider = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=Config.DEFAULT_STEPS)
-
-                    with gr.Row():
-                        with gr.Column():
-                            seed_slider = gr.Slider(label="Seed (-1 random)", minimum=-1, maximum=Config.MAX_SEED, step=1, value=-1)
-                            nums_slider = gr.Slider(label="Image Count", minimum=1, maximum=2, step=1, value=1)
-                        with gr.Column():
-                            lora_scale_slider = gr.Slider(label="LoRA Scale", minimum=0.1, maximum=2.0, step=0.1, value=Config.DEFAULT_LORA_SCALE)
-
-                    with gr.Row():
-                        with gr.Column():
-                            lora_add_text = gr.Textbox(label="Flux LoRA Path", lines=1, value=Config.DEFAULT_LORA)
-                        with gr.Column():
-                            lora_word_text = gr.Textbox(label="Flux LoRA Trigger Word", lines=1, value=Config.DEFAULT_TRIGGER_WORD)
-
-                    load_lora_btn = gr.Button(value="Load Custom LoRA", variant="secondary")
-
-            # Examples
-            gr.Examples(
-                examples=self.example_prompts,
-                inputs=[prompt_input, lora_word_text, lora_scale_slider],
-                outputs=[],
-                cache_examples=False,
-                examples_per_page=4
-            )
-
-            # Helper functions for UI status
-            def update_status_processing():
-                return "**⏳ Processing...**"
-
-            def update_status_done():
-                return "**✅ Done!**"
-
-            # Workflow for generate
-            generate_btn.click(
-                fn=update_status_processing, 
-                inputs=[], 
-                outputs=[status_markdown]
-            ).then(
-                fn=self.generator.generate,
-                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]
-            ).then(
-                fn=update_status_done,
-                inputs=[],
-                outputs=[status_markdown]
-            )
-
-            # Load LoRA
-            load_lora_btn.click(
-                fn=self.generator.load_lora,
-                inputs=[lora_add_text],
-                outputs=[lora_add_text]
-            )
-
-        return demo
+pipe = FluxPipeline.from_pretrained(
+    config["model_id"], torch_dtype=torch.bfloat16
+).to(device)
+pipe.load_lora_weights(config["default_lora"], weight_name=config["default_weight_name"])
 
 # -----------------------------------------------------------------------------
-# Main entry point
+# Function to load a new LoRA model
 # -----------------------------------------------------------------------------
-def main():
+def enable_lora(lora_add: str):
+    pipe.unload_lora_weights()
+    if not lora_add:
+        return gr.update(value="")
+    url = f"https://huggingface.co/{lora_add}/tree/main"
     try:
-        generator = FluxGenerator()
-        ui = FluxUI(generator)
-        demo = ui.build()
-        # Launch with default queue
-        demo.queue().launch()
+        pipe.load_lora_weights(lora_add)
+        return gr.update(label="LoRA Loaded Now")
     except Exception as e:
-        print(f"Application startup failed: {str(e)}")
-        with gr.Blocks() as error_demo:
-            gr.Markdown(f"# Error Starting Application\n\n{str(e)}")
-            gr.Markdown("Check logs for more details.")
-        error_demo.launch()
+        raise gr.Error(f"Failed to load {lora_add}: {e}")
 
-if __name__ == "__main__":
-    main()
+# -----------------------------------------------------------------------------
+# Function to generate an image from a prompt
+# -----------------------------------------------------------------------------
+@spaces.GPU()
+def generate_image(
+    prompt: str, lora_word: str, lora_scale: float = config["default_loRa_scale"],
+    width: int = config["default_width"], height: int = config["default_height"],
+    guidance_scale: float = config["default_guidance_scale"], steps: int = config["default_steps"],
+    seed: int = -1, nums: int = 1
+):
+    pipe.to(device)
+    seed = random.randint(0, config["max_seed"]) if seed == -1 else int(seed)
+    prompt_english = str(translator.translate(prompt, "English"))
+    full_prompt = f"{prompt_english} {lora_word}"
+    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, num_images_per_prompt=nums,
+        generator=generator, joint_attention_kwargs={"scale": lora_scale},
+    )
+    return result.images, seed
+
+# -----------------------------------------------------------------------------
+# Gradio UI
+# -----------------------------------------------------------------------------
+example_prompts = [
+    ["Medium-shot portrait, ohwx blue alien, wearing black techwear with a high collar, standing inside a futuristic VR showroom.", "ohwx", 0.9],
+    ["ohwx blue alien, wearing black techwear with a high collar, immersed in a digital cybernetic landscape.", "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.", "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.", "ohwx", 0.9]
+]
+
+with gr.Blocks(css=config["css"]) as demo:
+    gr.HTML("<h1><center>BR METAVERSO - Avatar Generator</center></h1>")
+    
+    processing_status = gr.Markdown("**🟢 Ready**", visible=True)  # Status indicator
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            gallery = gr.Gallery(label="Flux Generated Image", columns=1, preview=True, height=600)
+            prompt_input = gr.Textbox(label="Enter Your Prompt", lines=2, placeholder="Enter prompt...")
+            generate_btn = gr.Button(variant="primary")
+        with gr.Accordion("Advanced Options", open=True):
+            width_slider = gr.Slider(label="Width", minimum=512, maximum=1920, step=8, value=config["default_width"])
+            height_slider = gr.Slider(label="Height", minimum=512, maximum=1920, step=8, value=config["default_height"])
+            guidance_slider = gr.Slider(label="Guidance Scale", minimum=3.5, maximum=7, step=0.1, value=config["default_guidance_scale"])
+            steps_slider = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=config["default_steps"])
+            seed_slider = gr.Slider(label="Seed", minimum=-1, maximum=config["max_seed"], step=1, value=-1)
+            nums_slider = gr.Slider(label="Image Count", minimum=1, maximum=2, step=1, value=1)
+            lora_scale_slider = gr.Slider(label="LoRA Scale", minimum=0.1, maximum=2.0, step=0.1, value=config["default_loRa_scale"])
+            lora_add_text = gr.Textbox(label="Flux LoRA", lines=1, value=config["default_lora"])
+            lora_word_text = gr.Textbox(label="Flux LoRA Trigger Word", lines=1, value="ohwx")
+            load_lora_btn = gr.Button(value="Load LoRA", variant="secondary")
+    
+    gr.Examples(examples=example_prompts, inputs=[prompt_input, lora_word_text, lora_scale_slider], cache_examples=False, examples_per_page=4)
+    
+    # Ensuring processing status updates correctly
+    def update_status():
+        return "**⏳ Processing...**"
+    
+    generate_btn.click(fn=update_status, inputs=[], outputs=[processing_status]).then(
+        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]
+    ).then(
+        fn=lambda: "**✅ Done!**",
+        inputs=[],
+        outputs=[processing_status]
+    )
+    
+    load_lora_btn.click(fn=enable_lora, inputs=[lora_add_text], outputs=lora_add_text)
+    
+demo.queue().launch()