Update merged_app2.py

merged_app2.py

@@ -150,12 +150,9 @@ def download_models():
                 print(f"Error downloading {filename}: {e}")
 
 ensure_directories()
-
 download_models()
 
 
-
-
 hf_hub_download(repo_id="black-forest-labs/FLUX.1-Redux-dev", filename="flux1-redux-dev.safetensors", local_dir="models/style_models")
 hf_hub_download(repo_id="black-forest-labs/FLUX.1-Depth-dev", filename="flux1-depth-dev.safetensors", local_dir="models/diffusion_models")
 hf_hub_download(repo_id="Comfy-Org/sigclip_vision_384", filename="sigclip_vision_patch14_384.safetensors", local_dir="models/clip_vision")
@@ -165,7 +162,6 @@ hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="clip_l.sa
 t5_path = hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp16.safetensors", local_dir="models/text_encoders/t5")
 
 
-
 sd15_name = 'stablediffusionapi/realistic-vision-v51'
 tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer")
 text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder")
@@ -173,23 +169,10 @@ vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae")
 unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")
 
 
-
-# fill_pipe = FluxFillPipeline.from_single_file(
-#     "https://huggingface.co/SporkySporkness/FLUX.1-Fill-dev-GGUF/flux1-fill-dev-fp16-Q5_0-GGUF.gguf", 
-#     text_encoder= text_encoder,
-#     text_encoder_2 = t5_path,
-#     ignore_mismatched_sizes=True,
-#     low_cpu_mem_usage=False,
-#     torch_dtype=torch.bfloat16
-# ).to("cuda")
-
 from diffusers import FluxTransformer2DModel, FluxFillPipeline, GGUFQuantizationConfig
 from transformers import T5EncoderModel
 import torch
 
-# transformer = FluxTransformer2DModel.from_pretrained("AlekseyCalvin/FluxFillDev_fp8_Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16).to("cuda")
-# text_encoder_2 = T5EncoderModel.from_pretrained("AlekseyCalvin/FluxFillDev_fp8_Diffusers", subfolder="text_encoder_2", torch_dtype=torch.bfloat16).to("cuda")
-# fill_pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16).to("cuda")
 
 
 ckpt_path = (
@@ -209,8 +192,6 @@ fill_pipe = FluxFillPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
 )
 
-fill_pipe.enable_model_cpu_offload()
-
 
 try:
     import xformers
@@ -221,6 +202,11 @@ except ImportError:
     XFORMERS_AVAILABLE = False
     print("xformers not available - Using default attention")
 
+fill_pipe.enable_model_cpu_offload()
+fill_pipe.enable_vae_slicing()
+fill_pipe.enable_xformers_memory_efficient_attention()
+
+
 # Memory optimizations for RTX 2070
 torch.backends.cudnn.benchmark = True
 if torch.cuda.is_available():
@@ -505,6 +491,9 @@ pipe = prepare_pipeline(
     dtype=dtype,
 )
 
+pipe.enable_model_cpu_offload()
+pipe.enable_vae_slicing()
+pipe.enable_xformers_memory_efficient_attention()
 
 # Move models to device with consistent dtype
 text_encoder = text_encoder.to(device=device, dtype=dtype)
@@ -554,6 +543,10 @@ t2i_pipe = StableDiffusionPipeline(
     image_encoder=None
 )
 
+t2i_pipe.enable_model_cpu_offload()
+t2i_pipe.enable_vae_slicing()
+t2i_pipe.enable_xformers_memory_efficient_attention()
+
 i2i_pipe = StableDiffusionImg2ImgPipeline(
     vae=vae,
     text_encoder=text_encoder,
@@ -566,6 +559,9 @@ i2i_pipe = StableDiffusionImg2ImgPipeline(
     image_encoder=None
 )
 
+i2i_pipe.enable_model_cpu_offload()
+i2i_pipe.enable_vae_slicing()
+i2i_pipe.enable_xformers_memory_efficient_attention()
 
 @torch.inference_mode()
 def encode_prompt_inner(txt: str):
@@ -1054,7 +1050,7 @@ def use_orientation(selected_image:gr.SelectData):
 
 
 def generate_description(object_description,image, detail="high", max_tokens=250):
-    openai_api_key = os.getenv["OPENAI_API_KEY"]
+    openai_api_key = os.getenv("OPENAI_API_KEY")
     client = OpenAI(api_key=openai_api_key)
 
     if image is not None: