Update app.py

app.py

@@ -1,172 +1,172 @@
-import multiprocessing
+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
+
+# Instantiate the Accelerator
+accelerator = Accelerator()
+
+dtype = torch.bfloat16
+
+# Set environment variables for local path
+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
+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:
+        generated_text = text_generator(prompt, max_length=160, num_return_sequences=1, truncation=True)[0]['generated_text']
+        generated_description = re.sub(rf'{re.escape(prompt)}\s*', '', generated_text).strip()  # Remove the prompt from the generated text
+        return generated_description if generated_description else None
+    except Exception as e:
+        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()
+paused = False
+
+# Queue to store parsed descriptions
+parsed_descriptions_queue = deque()
+
+# Usage limits
+MAX_DESCRIPTIONS = 10
+MAX_IMAGES = 5
+
+@spaces.GPU
+def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_iterations=50):
+    global paused
+    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')
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    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:
+        if paused:
+            break
+
+        # 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.")
+            break
+
+        subject = random.choice(available_subjects)
+        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
+
+            # Parse and append descriptions every 3 iterations
+            if iteration_count % 3 == 0:
+                parsed_descriptions = parse_descriptions(clean_description)
+                parsed_descriptions_queue.extend(parsed_descriptions)
+
+        iteration_count += 1
+
+    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
+    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)
+    return images
 
 if __name__ == '__main__':
-    multiprocessing.set_start_method('spawn')
-
-    import spaces  # Import spaces at the very beginning
-    import os
-    import pandas as pd
-    import torch
-    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
-
-    # Instantiate the Accelerator
-    accelerator = Accelerator()
-
-    dtype = torch.bfloat16
-
-    # Set environment variables for local path
-    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
-    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:
-            generated_text = text_generator(prompt, max_length=160, num_return_sequences=1, truncation=True)[0]['generated_text']
-            generated_description = re.sub(rf'{re.escape(prompt)}\s*', '', generated_text).strip()  # Remove the prompt from the generated text
-            return generated_description if generated_description else None
-        except Exception as e:
-            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()
-    paused = False
-
-    # Queue to store parsed descriptions
-    parsed_descriptions_queue = deque()
-
-    # Usage limits
-    MAX_DESCRIPTIONS = 10
-    MAX_IMAGES = 5
-
-    @spaces.GPU
-    def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_iterations=50):
-        global paused
-        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')
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        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:
-            if paused:
-                break
-
-            # 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.")
-                break
-
-            subject = random.choice(available_subjects)
-            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
-
-                # Parse and append descriptions every 3 iterations
-                if iteration_count % 3 == 0:
-                    parsed_descriptions = parse_descriptions(clean_description)
-                    parsed_descriptions_queue.extend(parsed_descriptions)
-
-            iteration_count += 1
-
-        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
-        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)
-        return images
+    mp.set_start_method('spawn')
+    initialize_cuda()
 
     interface = gr.Interface(
         fn=combined_function,
@@ -175,3 +175,4 @@ if __name__ == '__main__':
     )
 
     interface.launch()
+