Update app.py

app.py

@@ -6,21 +6,12 @@ import sys
 import time
 import random
 import json
-from math import floor
 from typing import Any, Dict, List, Optional, Union
 
-# Local import for default LoRA list (if available)
-try:
-    from flux_app.lora import loras
-except ImportError:
-    loras = [
-        {"image": "placeholder.jpg", "title": "Placeholder LoRA", "repo": "placeholder/repo", "weights": None, "trigger_word": ""}
-    ]
-
 import torch
 import numpy as np
-import requests
 from PIL import Image
+import gradio as gr
 import spaces
 
 # Diffusers imports
@@ -32,12 +23,9 @@ from diffusers import (
 )
 from diffusers.utils import load_image
 
-# Hugging Face Hub
+# Hugging Face Hub imports
 from huggingface_hub import ModelCard, HfFileSystem
 
-# Gradio (UI)
-import gradio as gr
-
 ##############################
 # ===== config.py =====
 ##############################
@@ -50,28 +38,20 @@ MAX_SEED = 2**32 - 1
 ##############################
 # ===== utilities.py =====
 ##############################
-def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.16,
-):
+def calculate_shift(image_seq_len, base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.16):
     m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
     b = base_shift - m * base_seq_len
     mu = image_seq_len * m + b
     return mu
 
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
+def retrieve_timesteps(scheduler,
+                       num_inference_steps: Optional[int] = None,
+                       device: Optional[Union[str, torch.device]] = None,
+                       timesteps: Optional[List[int]] = None,
+                       sigmas: Optional[List[float]] = None,
+                       **kwargs):
     if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed.")
     if timesteps is not None:
         scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
         timesteps = scheduler.timesteps
@@ -86,11 +66,9 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 def load_image_from_path(image_path: str):
-    """Loads an image from a given file path."""
     return load_image(image_path)
 
 def randomize_seed_if_needed(randomize_seed: bool, seed: int, max_seed: int) -> int:
-    """Randomizes the seed if requested."""
     if randomize_seed:
         return random.randint(0, max_seed)
     return seed
@@ -98,40 +76,29 @@ def randomize_seed_if_needed(randomize_seed: bool, seed: int, max_seed: int) ->
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
-
     def __enter__(self):
         self.start_time = time.time()
         return self
-
     def __exit__(self, exc_type, exc_value, traceback):
         self.end_time = time.time()
-        self.elapsed_time = self.end_time - self.start_time
+        elapsed = self.end_time - self.start_time
         if self.activity_name:
-            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
+            print(f"Elapsed time for {self.activity_name}: {elapsed:.6f} seconds")
         else:
-            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
+            print(f"Elapsed time: {elapsed:.6f} seconds")
 
 ##############################
 # ===== enhance.py =====
 ##############################
 def generate(message, max_new_tokens=256, temperature=0.9, top_p=0.95, repetition_penalty=1.0):
-    """
-    Generates an enhanced prompt using a streaming Hugging Face API.
-    Enhances the given prompt under 100 words without changing its essence.
-    """
     SYSTEM_PROMPT = (
         "You are a prompt enhancer and your work is to enhance the given prompt under 100 words "
         "without changing the essence, only write the enhanced prompt and nothing else."
     )
     timestamp = time.time()
-    formatted_prompt = (
-        f"<s>[INST] SYSTEM: {SYSTEM_PROMPT} [/INST]"
-        f"[INST] {message} {timestamp} [/INST]"
-    )
-    
+    formatted_prompt = f"<s>[INST] SYSTEM: {SYSTEM_PROMPT} [/INST][INST] {message} {timestamp} [/INST]"
     api_url = "https://ruslanmv-hf-llm-api.hf.space/api/v1/chat/completions"
     headers = {"Content-Type": "application/json"}
-    
     payload = {
         "model": "mixtral-8x7b",
         "messages": [{"role": "user", "content": formatted_prompt}],
@@ -141,12 +108,10 @@ def generate(message, max_new_tokens=256, temperature=0.9, top_p=0.95, repetitio
         "use_cache": False,
         "stream": True
     }
-    
     try:
         response = requests.post(api_url, headers=headers, json=payload, stream=True)
         response.raise_for_status()
         full_output = ""
-        
         for line in response.iter_lines():
             if not line:
                 continue
@@ -172,32 +137,30 @@ def generate(message, max_new_tokens=256, temperature=0.9, top_p=0.95, repetitio
 ##############################
 # ===== lora_handling.py =====
 ##############################
-# A default list of LoRAs for the UI
+# Default LoRA list for initial UI setup
 loras = [
     {"image": "placeholder.jpg", "title": "Placeholder LoRA", "repo": "placeholder/repo", "weights": None, "trigger_word": ""}
 ]
 
 @torch.inference_mode()
-def flux_pipe_call_that_returns_an_iterable_of_images(
-    self,
-    prompt: Union[str, List[str]] = None,
-    prompt_2: Optional[Union[str, List[str]]] = None,
-    height: Optional[int] = None,
-    width: Optional[int] = None,
-    num_inference_steps: int = 28,
-    timesteps: List[int] = None,
-    guidance_scale: float = 3.5,
-    num_images_per_prompt: Optional[int] = 1,
-    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-    latents: Optional[torch.FloatTensor] = None,
-    prompt_embeds: Optional[torch.FloatTensor] = None,
-    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-    output_type: Optional[str] = "pil",
-    return_dict: bool = True,
-    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
-    max_sequence_length: int = 512,
-    good_vae: Optional[Any] = None,
-):
+def flux_pipe_call_that_returns_an_iterable_of_images(self,
+                                                       prompt: Union[str, List[str]] = None,
+                                                       prompt_2: Optional[Union[str, List[str]]] = None,
+                                                       height: Optional[int] = None,
+                                                       width: Optional[int] = None,
+                                                       num_inference_steps: int = 28,
+                                                       timesteps: List[int] = None,
+                                                       guidance_scale: float = 3.5,
+                                                       num_images_per_prompt: Optional[int] = 1,
+                                                       generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+                                                       latents: Optional[torch.FloatTensor] = None,
+                                                       prompt_embeds: Optional[torch.FloatTensor] = None,
+                                                       pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+                                                       output_type: Optional[str] = "pil",
+                                                       return_dict: bool = True,
+                                                       joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+                                                       max_sequence_length: int = 512,
+                                                       good_vae: Optional[Any] = None):
     height = height or self.default_sample_size * self.vae_scale_factor
     width = width or self.default_sample_size * self.vae_scale_factor
     
@@ -210,14 +173,11 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
         pooled_prompt_embeds=pooled_prompt_embeds,
         max_sequence_length=max_sequence_length,
     )
-
     self._guidance_scale = guidance_scale
     self._joint_attention_kwargs = joint_attention_kwargs
     self._interrupt = False
-
     batch_size = 1 if isinstance(prompt, str) else len(prompt)
     device = self._execution_device
-
     lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
     prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
         prompt=prompt,
@@ -229,7 +189,6 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
         max_sequence_length=max_sequence_length,
         lora_scale=lora_scale,
     )
-    
     num_channels_latents = self.transformer.config.in_channels // 4
     latents, latent_image_ids = self.prepare_latents(
         batch_size * num_images_per_prompt,
@@ -241,7 +200,6 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
         generator,
         latents,
     )
-    
     sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
     image_seq_len = latents.shape[1]
     mu = calculate_shift(
@@ -260,17 +218,13 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
         mu=mu,
     )
     self._num_timesteps = len(timesteps)
-
     guidance = (torch.full([1], guidance_scale, device=device, dtype=torch.float32)
                 .expand(latents.shape[0])
                 if self.transformer.config.guidance_embeds else None)
-
     for i, t in enumerate(timesteps):
         if self.interrupt:
             continue
-
         timestep = t.expand(latents.shape[0]).to(latents.dtype)
-
         noise_pred = self.transformer(
             hidden_states=latents,
             timestep=timestep / 1000,
@@ -282,14 +236,12 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
             joint_attention_kwargs=self.joint_attention_kwargs,
             return_dict=False,
         )[0]
-
         latents_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)
         latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor
         image = self.vae.decode(latents_for_image, return_dict=False)[0]
         yield self.image_processor.postprocess(image, output_type=output_type)[0]
         latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
         torch.cuda.empty_cache()
-        
     latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
     latents = (latents / good_vae.config.scaling_factor) + good_vae.config.shift_factor
     image = good_vae.decode(latents, return_dict=False)[0]
@@ -301,12 +253,10 @@ def get_huggingface_safetensors(link: str) -> tuple:
     split_link = link.split("/")
     if len(split_link) == 2:
         model_card = ModelCard.load(link)
-        base_model = model_card.data.get("base_model")
-        print(base_model)
-
-        if base_model not in ("black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"):
+        base_model_card = model_card.data.get("base_model")
+        print(base_model_card)
+        if base_model_card not in ("black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"):
             raise Exception("Flux LoRA Not Found!")
-
         image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
         trigger_word = model_card.data.get("instance_prompt", "")
         image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
@@ -319,12 +269,12 @@ def get_huggingface_safetensors(link: str) -> tuple:
                 if not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
                     image_elements = file.split("/")
                     image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-            return split_link[1], link, safetensors_name, trigger_word, image_url
         except Exception as e:
             print(e)
-            raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            raise Exception("Invalid LoRA repository")
+        return split_link[1], link, safetensors_name, trigger_word, image_url
     else:
-        raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+        raise Exception("Invalid LoRA link format")
 
 def check_custom_model(link: str) -> tuple:
     if link.startswith("https://"):
@@ -334,11 +284,8 @@ def check_custom_model(link: str) -> tuple:
     return get_huggingface_safetensors(link)
 
 def create_lora_card(title: str, repo: str, trigger_word: str, image: str) -> str:
-    trigger_word_info = (
-        f"Using: <code><b>{trigger_word}</b></code> as the trigger word"
-        if trigger_word
-        else "No trigger word found. If there's a trigger word, include it in your prompt"
-    )
+    trigger_word_info = (f"Using: <code><b>{trigger_word}</b></code> as the trigger word"
+                         if trigger_word else "No trigger word found. Include it in your prompt")
     return f'''
     <div class="custom_lora_card">
         <span>Loaded custom LoRA:</span>
@@ -352,14 +299,14 @@ def create_lora_card(title: str, repo: str, trigger_word: str, image: str) -> st
     </div>
     '''
 
-def add_custom_lora(custom_lora: str, loras_list: list) -> tuple:
+def add_custom_lora(custom_lora: str) -> tuple:
+    global loras
     if custom_lora:
         try:
             title, repo, path, trigger_word, image = check_custom_model(custom_lora)
             print(f"Loaded custom LoRA: {repo}")
             card = create_lora_card(title, repo, trigger_word, image)
-
-            existing_item_index = next((index for (index, item) in enumerate(loras_list) if item['repo'] == repo), None)
+            existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
             if existing_item_index is None:
                 new_item = {
                     "image": image,
@@ -369,11 +316,9 @@ def add_custom_lora(custom_lora: str, loras_list: list) -> tuple:
                     "trigger_word": trigger_word
                 }
                 print(new_item)
-                loras_list.append(new_item)
-                existing_item_index = len(loras_list) - 1
-
+                loras.append(new_item)
+                existing_item_index = len(loras) - 1
             return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
-
         except Exception as e:
             print(f"Error loading LoRA: {e}")
             return gr.update(visible=True, value="Invalid LoRA"), gr.update(visible=False), gr.update(), "", None, ""
@@ -386,7 +331,6 @@ def remove_custom_lora() -> tuple:
 def prepare_prompt(prompt: str, selected_index: Optional[int], loras_list: list) -> str:
     if selected_index is None:
         raise gr.Error("You must select a LoRA before proceeding.🧨")
-
     selected_lora = loras_list[selected_index]
     trigger_word = selected_lora.get("trigger_word")
     if trigger_word:
@@ -412,8 +356,8 @@ def load_lora_weights_into_pipeline(pipe_to_use, lora_path: str, weight_name: Op
         low_cpu_mem_usage=True
     )
 
-def update_selection(evt: gr.SelectData, width, height, loras_list):
-    selected_lora = loras_list[evt.index]
+def update_selection(evt: gr.SelectData, width, height) -> tuple:
+    selected_lora = loras[evt.index]
     new_placeholder = f"Type a prompt for {selected_lora['title']}"
     lora_repo = selected_lora["repo"]
     updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"
@@ -448,12 +392,9 @@ class ModelManager:
         self.initialize_models()
 
     def initialize_models(self):
-        """Initializes the diffusion pipelines and autoencoders."""
         self.taef1 = AutoencoderTiny.from_pretrained(TAEF1_MODEL, torch_dtype=DTYPE).to(DEVICE)
         self.good_vae = AutoencoderKL.from_pretrained(BASE_MODEL, subfolder="vae", torch_dtype=DTYPE).to(DEVICE)
-        # Optionally, pass use_auth_token=self.hf_token if needed.
-        self.pipe = DiffusionPipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE, vae=self.taef1)
-        self.pipe = self.pipe.to(DEVICE)
+        self.pipe = DiffusionPipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE, vae=self.taef1).to(DEVICE)
         self.pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
             BASE_MODEL,
             vae=self.good_vae,
@@ -464,14 +405,11 @@ class ModelManager:
             tokenizer_2=self.pipe.tokenizer_2,
             torch_dtype=DTYPE,
         ).to(DEVICE)
-        # Instead of binding to the instance (which fails due to __slots__),
-        # bind the custom method to the pipeline’s class.
+        # Bind custom LoRA method to the pipeline class (to avoid __slots__ issues)
         self.pipe.__class__.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images
     
-    @spaces.GPU(duration=100)
+    @spaces.GPU(duration=100) 
     def generate_image(self, prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
-        """Generates an image using the text-to-image pipeline."""
-        self.pipe.to(DEVICE)
         generator = torch.Generator(device=DEVICE).manual_seed(seed)
         with calculateDuration("Generating image"):
             for img in self.pipe.flux_pipe_call_that_returns_an_iterable_of_images(
@@ -488,9 +426,7 @@ class ModelManager:
                 yield img
 
     def generate_image_to_image(self, prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
-        """Generates an image using the image-to-image pipeline."""
         generator = torch.Generator(device=DEVICE).manual_seed(seed)
-        self.pipe_i2i.to(DEVICE)
         image_input = load_image_from_path(image_input_path)
         with calculateDuration("Generating image to image"):
             final_image = self.pipe_i2i(
@@ -513,12 +449,11 @@ class ModelManager:
 class Frontend:
     def __init__(self, model_manager: ModelManager):
         self.model_manager = model_manager
-        self.loras = loras  # Use the default LoRA list defined above.
+        self.loras = loras
         self.load_initial_loras()
         self.css = self.define_css()
 
     def define_css(self):
-        # Clean and professional CSS styling.
         return '''
         /* Title Styling */
         #title {
@@ -587,18 +522,14 @@ class Frontend:
             self.loras = loras_list
         except ImportError:
             print("Warning: lora.py not found, using placeholder LoRAs.")
-            pass
 
     @spaces.GPU(duration=100)
     def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index,
                  randomize_seed, seed, width, height, lora_scale, use_enhancer,
                  progress=gr.Progress(track_tqdm=True)):
         seed = randomize_seed_if_needed(randomize_seed, seed, MAX_SEED)
-        # Prepare the prompt using the selected LoRA trigger word.
         prompt_mash = prepare_prompt(prompt, selected_index, self.loras)
         enhanced_text = ""
-        
-        # Optionally enhance the prompt.
         if use_enhancer:
             for enhanced_chunk in generate(prompt_mash):
                 enhanced_text = enhanced_chunk
@@ -606,12 +537,10 @@ class Frontend:
             prompt_mash = enhanced_text
         else:
             enhanced_text = ""
-        
         selected_lora = self.loras[selected_index]
         unload_lora_weights(self.model_manager.pipe, self.model_manager.pipe_i2i)
         pipe_to_use = self.model_manager.pipe_i2i if image_input is not None else self.model_manager.pipe
         load_lora_weights_into_pipeline(pipe_to_use, selected_lora["repo"], selected_lora.get("weights"))
-
         if image_input is not None:
             final_image = self.model_manager.generate_image_to_image(
                 prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed
@@ -630,12 +559,8 @@ class Frontend:
 
     def create_ui(self):
         with gr.Blocks(theme=gr.themes.Base(), css=self.css, title="Flux LoRA Generation") as app:
-            title = gr.HTML(
-                """<h1>Flux LoRA Generation</h1>""",
-                elem_id="title",
-            )
+            title = gr.HTML("<h1>Flux LoRA Generation</h1>", elem_id="title")
             selected_index = gr.State(None)
-
             with gr.Row():
                 with gr.Column(scale=3):
                     prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Choose the LoRA and type the prompt")
@@ -660,7 +585,6 @@ class Frontend:
                 with gr.Column():
                     progress_bar = gr.Markdown(elem_id="progress", visible=False)
                     result = gr.Image(label="Generated Image")
-
             with gr.Row():
                 with gr.Accordion("Advanced Settings", open=False):
                     with gr.Row():
@@ -681,44 +605,37 @@ class Frontend:
                             use_enhancer = gr.Checkbox(value=False, label="Use Prompt Enhancer")
                             show_enhanced_prompt = gr.Checkbox(value=False, label="Display Enhanced Prompt")
                     enhanced_prompt_box = gr.Textbox(label="Enhanced Prompt", visible=False)
-
             gallery.select(
                 update_selection,
-                inputs=[width, height, gr.State(self.loras)],
+                inputs=[width, height],
                 outputs=[prompt, selected_info, selected_index, width, height]
             )
             custom_lora.input(
                 add_custom_lora,
-                inputs=[custom_lora, gr.State(self.loras)],
+                inputs=[custom_lora],
                 outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
             )
             custom_lora_button.click(
                 remove_custom_lora,
                 outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
             )
-
             show_enhanced_prompt.change(fn=lambda show: gr.update(visible=show),
                                         inputs=show_enhanced_prompt,
                                         outputs=enhanced_prompt_box)
-
             gr.on(
                 triggers=[generate_button.click, prompt.submit],
                 fn=self.run_lora,
-                inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index,
-                        randomize_seed, seed, width, height, lora_scale, use_enhancer],
+                inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, use_enhancer],
                 outputs=[result, seed, progress_bar, enhanced_prompt_box]
             )
-
             with gr.Row():
                 gr.HTML("<div style='text-align:center; font-size:0.9em; margin-top:20px;'>Credits: <a href='https://ruslanmv.com' target='_blank'>ruslanmv.com</a></div>")
-            
             return app
 
 ##############################
 # ===== Main app.py =====
 ##############################
 if __name__ == "__main__":
-    # Get the Hugging Face token from the environment.
     hf_token = os.environ.get("HF_TOKEN")
     if not hf_token:
         raise ValueError("Hugging Face token (HF_TOKEN) not found in environment variables. Please set it.")
@@ -726,5 +643,4 @@ if __name__ == "__main__":
     frontend = Frontend(model_manager)
     app = frontend.create_ui()
     app.queue()
-    # Set share=True to create a public link if desired.
     app.launch(share=False, debug=True)