Update app.py

app.py

@@ -1,270 +1,89 @@
 import os
-import random
-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"
-
-
-# -----------------------------------------------------------------------------
-# 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 = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using {self.device.upper()}")
-        
-        # Initialize pipeline
-        self.pipe = self._initialize_pipeline()
-        
-    def _initialize_pipeline(self):
-        """Initialize the Flux pipeline and load default LoRA"""
-        try:
-            pipe = FluxPipeline.from_pretrained(
-                Config.MODEL_ID, torch_dtype=torch.bfloat16
-            ).to(self.device)
-            pipe.load_lora_weights(Config.DEFAULT_LORA, weight_name=Config.DEFAULT_WEIGHT_NAME)
-            return pipe
-        except Exception as e:
-            print(f"Error initializing pipeline: {e}")
-            raise
+import json
+import torch
+import random
+import time
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
+from diffusers.utils import load_image
+from huggingface_hub import ModelCard
 
-    def load_lora(self, lora_path):
-        """Load a new LoRA model"""
-        try:
-            self.pipe.unload_lora_weights()
-            if not lora_path:
-                return gr.update(value="")
-                
-            self.pipe.load_lora_weights(lora_path)
-            return gr.update(label="LoRA Loaded Successfully")
-        except Exception as e:
-            error_msg = f"Failed to load LoRA from {lora_path}: {e}"
-            print(error_msg)
-            raise gr.Error(error_msg)
+# -------------------------------------------------------------------------
+# CONFIGURAÇÃO GERAL
+# -------------------------------------------------------------------------
+CONFIG = {
+    "base_model": "black-forest-labs/FLUX.1-dev",
+    "dtype": torch.float16,  # Substituído por torch.float16 para economizar VRAM
+    "device": "cuda" if torch.cuda.is_available() else "cpu",
+    "max_seed": 2**32 - 1
+}
 
-    @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"""
-        try:
-            # Ensure the pipe is on the correct device
-            self.pipe.to(self.device)
-            
-            # Handle seed
-            seed = random.randint(0, Config.MAX_SEED) if seed == -1 else int(seed)
-            generator = torch.Generator().manual_seed(seed)
-            
-            # Translate prompt if not in English
-            prompt_english = str(self.translator.translate(prompt, "English"))
-            full_prompt = f"{prompt_english} {lora_word}"
-            
-            # Generate image
-            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},
-            )
-            
-            return result.images, seed
-            
-        except Exception as e:
-            error_msg = f"Image generation failed: {e}"
-            print(error_msg)
-            raise gr.Error(error_msg)
+# Limpa cache da GPU para evitar erro de falta de memória
+torch.cuda.empty_cache()
+torch.cuda.ipc_collect()
 
+# -------------------------------------------------------------------------
+# CARREGANDO O MODELO BASE
+# -------------------------------------------------------------------------
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=CONFIG["dtype"]).to(CONFIG["device"])
+good_vae = AutoencoderKL.from_pretrained(CONFIG["base_model"], subfolder="vae", torch_dtype=CONFIG["dtype"]).to(CONFIG["device"])
 
-# -----------------------------------------------------------------------------
-# UI Builder class
-# -----------------------------------------------------------------------------
-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):
-        """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)
-            
-            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..."
-                    )
-                    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 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
-                            )
-                        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 section
-            gr.Examples(
-                examples=self.example_prompts,
-                inputs=[prompt_input, lora_word_text, lora_scale_slider],
-                cache_examples=False,
-                examples_per_page=4
-            )
-            
-            # Wire up the event handlers
-            # Status update functions
-            def update_status_processing():
-                return "**⏳ Processing...**"
-                
-            def update_status_done():
-                return "**✅ Done!**"
+pipe = DiffusionPipeline.from_pretrained(
+    CONFIG["base_model"],
+    torch_dtype=CONFIG["dtype"],
+    vae=taef1,
+    low_cpu_mem_usage=True  # Economiza memória na GPU
+).to(CONFIG["device"])
 
-            # Generate button click workflow
-            generate_btn.click(
-                fn=update_status_processing,
-                inputs=[],
-                outputs=[processing_status]
-            ).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=[processing_status]
-            )
-            
-            # Load LoRA button click workflow
-            load_lora_btn.click(
-                fn=self.generator.load_lora,
-                inputs=[lora_add_text],
-                outputs=[lora_add_text]
-            )
-            
-        return demo
+pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
+    CONFIG["base_model"],
+    vae=good_vae,
+    transformer=pipe.transformer,
+    text_encoder=pipe.text_encoder,
+    tokenizer=pipe.tokenizer,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer_2=pipe.tokenizer_2,
+    torch_dtype=CONFIG["dtype"]
+)
 
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
-# -----------------------------------------------------------------------------
-# Main application
-# -----------------------------------------------------------------------------
-def main():
-    # Initialize generator
-    generator = FluxGenerator()
+# -------------------------------------------------------------------------
+# FUNÇÃO PARA GERAR IMAGEM
+# -------------------------------------------------------------------------
+def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, progress):
+    pipe.to(CONFIG["device"])  # Garante que o modelo está na GPU
+    generator = torch.Generator(device=CONFIG["device"]).manual_seed(seed)
+    
+    # Medir tempo de geração
+    start_time = time.time()
     
-    # Build and launch UI
-    ui = FluxUI(generator)
-    demo = ui.build()
-    demo.queue().launch()
+    # Gerando a imagem
+    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+        prompt=prompt,
+        num_inference_steps=steps,
+        guidance_scale=cfg_scale,
+        width=width,
+        height=height,
+        generator=generator,
+        joint_attention_kwargs={"scale": lora_scale},
+        output_type="pil",
+        good_vae=good_vae,
+    ):
+        end_time = time.time()
+        print(f"Tempo de geração: {end_time - start_time:.2f}s")
+        yield img
 
+# -------------------------------------------------------------------------
+# INTERFACE GRADIO
+# -------------------------------------------------------------------------
+with gr.Blocks(theme=gr.themes.Soft()) as app:
+    gr.Markdown("# FLUX Avatar Generator")
 
-if __name__ == "__main__":
-    main()
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", placeholder="Descreva sua imagem...")
+            steps = gr.Slider(1, 50, 25, step=1, label="Passos")
+            cfg_scale = gr.Slider(1, 20, 3.5, step=0.5, label="Escala CFG")
+            width = gr.Slider(512, 1536, 896, step=64, label="Largura