Update app.py

app.py

@@ -1,18 +1,9 @@
-import spaces # beginn
 import torch.multiprocessing as mp
 import torch
 import os
-import pandas as pd
-import gc
 import re
 import random
-from tqdm.auto import tqdm
 from collections import deque
-from optimum.quanto import freeze, qfloat8, quantize
-from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL
-from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
-from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 from accelerate import Accelerator
@@ -30,40 +21,56 @@ dtype = torch.bfloat16
 os.environ['FLUX_DEV'] = '.'
 os.environ['AE'] = '.'
 
-bfl_repo = 'black-forest-labs/FLUX.1-schnell'
-revision = 'refs/pr/1'
-
-scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder='scheduler', revision=revision)
-text_encoder = CLIPTextModel.from_pretrained('openai/clip-vit-large-patch14', torch_dtype=dtype)
-tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14', torch_dtype=dtype)
-text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder='text_encoder_2', torch_dtype=dtype, revision=revision)
-tokenizer_2 = T5TokenizerFast.from_pretrained(bfl_repo, subfolder='tokenizer_2', torch_dtype=dtype, revision=revision)
-vae = AutoencoderKL.from_pretrained(bfl_repo, subfolder='vae', torch_dtype=dtype, revision=revision)
-transformer = FluxTransformer2DModel.from_pretrained(bfl_repo, subfolder='transformer', torch_dtype=dtype, revision=revision)
-
-quantize(transformer, weights=qfloat8)
-freeze(transformer)
-quantize(text_encoder_2, weights=qfloat8)
-freeze(text_encoder_2)
-
-pipe = FluxPipeline(
-    scheduler=scheduler,
-    text_encoder=text_encoder,
-    tokenizer=tokenizer,
-    text_encoder_2=None,
-    tokenizer_2=tokenizer_2,
-    vae=vae,
-    transformer=None,
-)
-pipe.text_encoder_2 = text_encoder_2
-pipe.transformer = transformer
-pipe.enable_model_cpu_offload()
-
-# Create a directory to save the generated images
-output_dir = 'generated_images'
-os.makedirs(output_dir, exist_ok=True)
-
-# Function to generate a detailed visual description prompt
+# Seed words pool
+seed_words = []
+
+used_words = set()
+
+# Queue to store parsed descriptions
+parsed_descriptions_queue = deque()
+
+# Usage limits
+MAX_DESCRIPTIONS = 30
+MAX_IMAGES = 12
+
+def initialize_diffusers():
+    from optimum.quanto import freeze, qfloat8, quantize
+    from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL
+    from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
+    from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
+    from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+    bfl_repo = 'black-forest-labs/FLUX.1-schnell'
+    revision = 'refs/pr/1'
+
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder='scheduler', revision=revision)
+    text_encoder = CLIPTextModel.from_pretrained('openai/clip-vit-large-patch14', torch_dtype=dtype)
+    tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14', torch_dtype=dtype)
+    text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder='text_encoder_2', torch_dtype=dtype, revision=revision)
+    tokenizer_2 = T5TokenizerFast.from_pretrained(bfl_repo, subfolder='tokenizer_2', torch_dtype=dtype, revision=revision)
+    vae = AutoencoderKL.from_pretrained(bfl_repo, subfolder='vae', torch_dtype=dtype, revision=revision)
+    transformer = FluxTransformer2DModel.from_pretrained(bfl_repo, subfolder='transformer', torch_dtype=dtype, revision=revision)
+
+    quantize(transformer, weights=qfloat8)
+    freeze(transformer)
+    quantize(text_encoder_2, weights=qfloat8)
+    freeze(text_encoder_2)
+
+    pipe = FluxPipeline(
+        scheduler=scheduler,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        text_encoder_2=None,
+        tokenizer_2=tokenizer_2,
+        vae=vae,
+        transformer=None,
+    )
+    pipe.text_encoder_2 = text_encoder_2
+    pipe.transformer = transformer
+    pipe.enable_model_cpu_offload()
+
+    return pipe
+
 def generate_description_prompt(subject, user_prompt, text_generator):
     prompt = f"write concise vivid visual description enclosed in brackets like [ <description> ] less than 100 words of {user_prompt} different from {subject}. "
     try:
@@ -74,33 +81,17 @@ def generate_description_prompt(subject, user_prompt, text_generator):
         print(f"Error generating description for subject '{subject}': {e}")
         return None
 
-# Function to parse descriptions from a given text
 def parse_descriptions(text):
-    # Find all descriptions enclosed in brackets
     descriptions = re.findall(r'\[([^\[\]]+)\]', text)
-    # Filter descriptions with at least 3 words
     descriptions = [desc.strip() for desc in descriptions if len(desc.split()) >= 3]
     return descriptions
 
-# Seed words pool
-seed_words = []
-
-used_words = set()
-
-# Queue to store parsed descriptions
-parsed_descriptions_queue = deque()
-
-# Usage limits
-MAX_DESCRIPTIONS = 30
-MAX_IMAGES = 12
-
 @spaces.GPU
 def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_iterations=50):
     descriptions = []
     description_queue = deque()
     iteration_count = 0
 
-    # Initialize the text generation pipeline with 16-bit precision
     print("Initializing the text generation pipeline with 16-bit precision...")
     model_name = 'meta-llama/Meta-Llama-3.1-8B-Instruct'
     model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map='auto')
@@ -108,11 +99,9 @@ def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_ite
     text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
     print("Text generation pipeline initialized with 16-bit precision.")
 
-    # Populate the seed_words array with user input
     seed_words.extend(re.findall(r'"(.*?)"', seed_words_input))
 
     while iteration_count < max_iterations and len(parsed_descriptions_queue) < MAX_DESCRIPTIONS:
-        # Select a subject that has not been used
         available_subjects = [word for word in seed_words if word not in used_words]
         if not available_subjects:
             print("No more available subjects to use.")
@@ -122,18 +111,14 @@ def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_ite
         generated_description = generate_description_prompt(subject, user_prompt, text_generator)
         
         if generated_description:
-            # Remove any offending symbols
             clean_description = generated_description.encode('ascii', 'ignore').decode('ascii')
             description_queue.append({'subject': subject, 'description': clean_description})
 
-            # Print the generated description to the command line
             print(f"Generated description for subject '{subject}': {clean_description}")
 
-            # Update used words and seed words
             used_words.add(subject)
-            seed_words.append(clean_description)  # Add the generated description to the seed bank array
+            seed_words.append(clean_description)
 
-            # Parse and append descriptions every 3 iterations
             if iteration_count % 3 == 0:
                 parsed_descriptions = parse_descriptions(clean_description)
                 parsed_descriptions_queue.extend(parsed_descriptions)
@@ -143,28 +128,25 @@ def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_ite
     return list(parsed_descriptions_queue)
 
 @spaces.GPU(duration=120)
-def generate_images(parsed_descriptions):
-    # If there are fewer than MAX_IMAGES descriptions, use whatever is available
+def generate_images(parsed_descriptions, pipe):
     if len(parsed_descriptions) < MAX_IMAGES:
         prompts = parsed_descriptions
     else:
         prompts = [parsed_descriptions.pop(0) for _ in range(MAX_IMAGES)]
 
-    # Generate images from the parsed descriptions
     images = []
     for prompt in prompts:
         images.extend(pipe(prompt, num_images=1).images)
 
     return images
 
-# Create Gradio Interface
 def combined_function(user_prompt, seed_words_input):
     parsed_descriptions = generate_descriptions(user_prompt, seed_words_input)
-    images = generate_images(parsed_descriptions)
+    pipe = initialize_diffusers()
+    images = generate_images(parsed_descriptions, pipe)
     return images
 
 if __name__ == '__main__':
-    # Ensure CUDA is initialized correctly
     torch.cuda.init()
 
     interface = gr.Interface(
@@ -173,4 +155,4 @@ if __name__ == '__main__':
         outputs=gr.Gallery()
     )
 
-    interface.launch()
+    interface.launch()