Update app.py

app.py

@@ -41,35 +41,45 @@ model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float
 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()
+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
 
+pipe = initialize_diffusers()
 print("Models and checkpoints preloaded.")
 
 def generate_description_prompt(subject, user_prompt, text_generator):
@@ -122,7 +132,7 @@ 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, max_iterations=2):  # Set max_iterations to 1
+def generate_images(parsed_descriptions, max_iterations=1):  # Set max_iterations to 1
     if len(parsed_descriptions) < MAX_IMAGES:
         prompts = parsed_descriptions
     else:
@@ -152,4 +162,4 @@ if __name__ == '__main__':
         allow_flagging='never'  # Disable flagging
     )
 
-    interface.launch(share=True)
+    interface.launch(share=True)