Update app.py

app.py

@@ -1,168 +1,177 @@
-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()
-
-@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:
-        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 13 descriptions, use whatever is available
-    if len(parsed_descriptions) < 13:
-        prompts = parsed_descriptions
-    else:
-        prompts = [parsed_descriptions.pop(0) for _ in range(13)]
-
-    # 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
-
-interface = gr.Interface(
-    fn=combined_function,
-    inputs=[gr.Textbox(lines=2, placeholder="Enter a prompt for descriptions..."), gr.Textbox(lines=2, placeholder='Enter seed words in quotes, e.g., "cat", "dog", "sunset"...')],
-    outputs=gr.Gallery()
-)
-
-interface.launch()
+import multiprocessing
+
+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
+
+    interface = gr.Interface(
+        fn=combined_function,
+        inputs=[gr.Textbox(lines=2, placeholder="Enter a prompt for descriptions..."), gr.Textbox(lines=2, placeholder='Enter seed words in quotes, e.g., "cat", "dog", "sunset"...')],
+        outputs=gr.Gallery()
+    )
+
+    interface.launch()