Update app.py

app.py

@@ -3,7 +3,6 @@ import random
 import torch
 import numpy as np
 import gradio as gr
-import spaces
 from diffusers import FluxPipeline
 from translatepy import Translator
 
@@ -27,77 +26,64 @@ class Config:
     ENABLE_SEQUENTIAL_CPU_OFFLOAD = True
     ENABLE_ATTENTION_SLICING = "max"
 
-
 # -----------------------------------------------------------------------------
 # FluxGenerator class to handle image generation
 # -----------------------------------------------------------------------------
 class FluxGenerator:
     def __init__(self):
-        # Environment setup
         os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
         self.translator = Translator()
         self.device = self._get_optimal_device()
         print(f"Using {self.device.upper()}")
-        
-        # Initialize pipeline
+
         self.pipe = None
         self._initialize_pipeline()
-        
+
     def _get_optimal_device(self):
         """Determine the optimal device based on available resources"""
         if torch.cuda.is_available():
-            # Check GPU memory
             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, using CPU with GPU acceleration")
-                    return "cuda"  # Still use CUDA but will apply memory optimizations
+                    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...")
-            # Use more memory-efficient settings
             pipe_kwargs = {
                 "torch_dtype": torch.bfloat16 if self.device == "cuda" else torch.float32,
             }
-            
-            # Initialize the pipeline
-            self.pipe = FluxPipeline.from_pretrained(
-                Config.MODEL_ID,
-                **pipe_kwargs
-            )
-            
+            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 using CPU offload
-                self.pipe = self.pipe.to(self.device)
-            
-            # Load default LoRA
+                # 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")
-            return self.pipe
-            
         except Exception as e:
             error_msg = f"Error initializing pipeline: {str(e)}"
             print(error_msg)
@@ -106,18 +92,17 @@ class FluxGenerator:
     def load_lora(self, lora_path):
         """Load a new LoRA model"""
         try:
-            print(f"Unloading previous LoRA weights...")
+            print("Unloading previous LoRA weights...")
             self.pipe.unload_lora_weights()
-            
+
             if not lora_path:
-                print("No LoRA path provided, skipping LoRA loading")
+                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")
+            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)
@@ -129,50 +114,43 @@ class FluxGenerator:
             try:
                 print("Clearing CUDA memory cache...")
                 torch.cuda.empty_cache()
-                if hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
+                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)}")
 
-    @spaces.GPU()
     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"""
+                 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}'")
-            
-            # Clear memory before generation
             self._clear_memory()
-            
-            # Ensure we're using the right device
+
             if not Config.ENABLE_SEQUENTIAL_CPU_OFFLOAD:
                 print(f"Moving model to {self.device}")
                 self.pipe.to(self.device)
-            
-            # Handle seed
+
             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)
-            
-            # Translate prompt if not in English
+
             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}'")
-            
-            # Lower resolution if on limited memory
-            if self.device == "cuda" and torch.cuda.get_device_properties(0).total_memory < 8 * 1024 * 1024 * 1024:
+
+            # 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
-                # Scale down to 85% if memory is tight
                 width = int(width * 0.85)
                 height = int(height * 0.85)
-                print(f"Limited memory detected. Scaling down resolution from {original_width}x{original_height} to {width}x{height}")
-            
-            # Generate with autocast for memory efficiency
+                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.cuda.amp.autocast(enabled=self.device == "cuda"):
+            with torch.autocast("cuda", enabled=(self.device == "cuda")):
                 result = self.pipe(
                     prompt=full_prompt,
                     height=height,
@@ -184,15 +162,12 @@ class FluxGenerator:
                     generator=generator,
                     joint_attention_kwargs={"scale": lora_scale},
                 )
-            
-            print("Generation complete, returning images")
-            self._clear_memory()  # Clear memory after generation
+            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)
-            # Clear memory after error
             self._clear_memory()
             raise gr.Error(error_msg)
 
@@ -209,135 +184,68 @@ class FluxUI:
             ["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):
-        """Build and return the Gradio interface"""
         with gr.Blocks(css=Config.CSS) as demo:
             gr.HTML("<h1><center>BR METAVERSO - Avatar Generator</center></h1>")
-            
-            # Status indicator
-            processing_status = gr.Markdown("**🟢 Ready**", visible=True)
-            
+
+            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="Enter prompt for your avatar..."
+                        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
-                            )
+                            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
-                            )
-                    
+                            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 for 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
-                            )
+                            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
-                            )
-                            
+                            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
-                            )
+                            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
-                            )
-                    
+                            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")
-                    
-                    # Memory optimization checkbox
-                    with gr.Row():
-                        memory_efficient = gr.Checkbox(
-                            label="Enable Memory Optimizations",
-                            value=True,
-                            info="Reduces memory usage but may increase generation time"
-                        )
-            
-            # Examples section
+
+            # Examples
             gr.Examples(
                 examples=self.example_prompts,
                 inputs=[prompt_input, lora_word_text, lora_scale_slider],
+                outputs=[],
                 cache_examples=False,
                 examples_per_page=4
             )
-            
-            # Wire up the event handlers
-            # Status update functions
+
+            # Helper functions for UI status
             def update_status_processing():
                 return "**⏳ Processing...**"
-                
+
             def update_status_done():
                 return "**✅ Done!**"
-                
-            def update_memory_settings(enable_memory_opt):
-                global Config
-                Config.ENABLE_MEMORY_EFFICIENT_ATTENTION = enable_memory_opt
-                Config.ENABLE_SEQUENTIAL_CPU_OFFLOAD = enable_memory_opt
-                Config.ENABLE_ATTENTION_SLICING = "max" if enable_memory_opt else None
-                return gr.update()
-
-            # Generate button click workflow
+
+            # Workflow for generate
             generate_btn.click(
-                fn=update_status_processing,
-                inputs=[],
-                outputs=[processing_status]
+                fn=update_status_processing, 
+                inputs=[], 
+                outputs=[status_markdown]
             ).then(
                 fn=self.generator.generate,
                 inputs=[
@@ -349,45 +257,34 @@ class FluxUI:
             ).then(
                 fn=update_status_done,
                 inputs=[],
-                outputs=[processing_status]
+                outputs=[status_markdown]
             )
-            
-            # Load LoRA button click workflow
+
+            # Load LoRA
             load_lora_btn.click(
                 fn=self.generator.load_lora,
                 inputs=[lora_add_text],
                 outputs=[lora_add_text]
             )
-            
-            # Memory optimization checkbox event
-            memory_efficient.change(
-                fn=update_memory_settings,
-                inputs=[memory_efficient],
-                outputs=[]
-            )
-            
-        return demo
 
+        return demo
 
 # -----------------------------------------------------------------------------
-# Main application
+# Main entry point
 # -----------------------------------------------------------------------------
 def main():
     try:
-        # Create a generator with memory optimizations
         generator = FluxGenerator()
-        
-        # Build and launch UI
         ui = FluxUI(generator)
         demo = ui.build()
-        
-        # Launch with low cache size to prevent memory issues
-        demo.queue(max_size=1).launch(share=False)
-        
+        # Launch with default queue
+        demo.queue().launch()
     except Exception as e:
         print(f"Application startup failed: {str(e)}")
-        # Show error in UI if possible
         with gr.Blocks() as error_demo:
-            gr.Markdown(f"# Error Starting Application\n\n{str(e)}\n\nPlease check the logs for more details.")
-            gr.Markdown("This might be due to memory limitations or GPU compatibility issues.")
-            error_demo.launch()
+            gr.Markdown(f"# Error Starting Application\n\n{str(e)}")
+            gr.Markdown("Check logs for more details.")
+        error_demo.launch()
+
+if __name__ == "__main__":
+    main()