load model before gpu spaces invoke

app.py

@@ -34,43 +34,43 @@ parsed_descriptions_queue = deque()
 MAX_DESCRIPTIONS = 30
 MAX_IMAGES = 1  # Generate only 1 image
 
-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
+# Preload models and checkpoints
+print("Preloading models and checkpoints...")
+model_name = 'NousResearch/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)
+
+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_clip = 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_clip,
+    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()
+
+print("Models and checkpoints preloaded.")
 
 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}. "
@@ -93,13 +93,6 @@ def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_ite
     description_queue = deque()
     iteration_count = 0
 
-    print("Initializing the text generation pipeline with 16-bit precision...")
-    model_name = 'NousResearch/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.")
-
     seed_words.extend(re.findall(r'"(.*?)"', seed_words_input))
 
     for _ in range(2):  # Perform two iterations
@@ -128,10 +121,8 @@ def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_ite
 
     return list(parsed_descriptions_queue)
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def generate_images(parsed_descriptions, max_iterations=2):  # Set max_iterations to 1
-    pipe = initialize_diffusers()
-
     if len(parsed_descriptions) < MAX_IMAGES:
         prompts = parsed_descriptions
     else:
@@ -161,4 +152,4 @@ if __name__ == '__main__':
         allow_flagging='never'  # Disable flagging
     )
 
-    interface.launch(share=True)
+    interface.launch(share=True)