Update app.py

app.py

@@ -1,3 +1,46 @@
+import subprocess
+import sys
+
+# Install/upgrade required packages with specific versions
+def install_packages():
+    packages = [
+        "transformers>=4.46.0",
+        "diffusers>=0.31.0",
+        "accelerate>=0.26.0",
+        "huggingface-hub>=0.23.0"
+    ]
+    
+    for package in packages:
+        subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", package], check=True)
+
+# Run installation before other imports
+try:
+    install_packages()
+except Exception as e:
+    print(f"Warning: Could not auto-install packages: {e}")
+
+# Try to install flash-attn with a timeout
+try:
+    print("Attempting to install flash-attn...")
+    result = subprocess.run(
+        'pip install flash-attn --no-build-isolation', 
+        env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, 
+        shell=True,
+        timeout=120,  # 2 minute timeout
+        capture_output=True,
+        text=True
+    )
+    if result.returncode == 0:
+        print("Flash-attn installed successfully")
+    else:
+        print(f"Flash-attn installation failed: {result.stderr}")
+        print("Continuing without flash-attn...")
+except subprocess.TimeoutExpired:
+    print("Flash-attn installation timed out - continuing without it")
+except Exception as e:
+    print(f"Flash-attn installation error: {e}")
+    print("Continuing without flash-attn...")
+
 import spaces
 import argparse
 import os
@@ -13,12 +56,21 @@ from diffusers import FluxPipeline
 from diffusers.pipelines.stable_diffusion import safety_checker
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM
-import subprocess
 
-# Flash Attention 설치
-subprocess.run('pip install flash-attn --no-build-isolation', 
-              env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, 
-              shell=True)
+# Try to use efficient attention mechanisms
+ATTN_METHOD = None
+try:
+    import xformers
+    print("Using xformers for efficient attention")
+    ATTN_METHOD = "xformers"
+except ImportError:
+    try:
+        import flash_attn
+        print("Using flash attention")
+        ATTN_METHOD = "flash_attn"
+    except ImportError:
+        print("No efficient attention method available, using default")
+        ATTN_METHOD = "default"
 
 # Setup and initialization code
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
@@ -31,6 +83,7 @@ os.environ["HF_HOME"] = cache_path
 torch.backends.cuda.matmul.allow_tf32 = True
 
 # Florence 모델 초기화
+print("Initializing Florence models...")
 florence_models = {
     'gokaygokay/Florence-2-Flux-Large': AutoModelForCausalLM.from_pretrained(
         'gokaygokay/Florence-2-Flux-Large', 
@@ -80,10 +133,13 @@ class timer:
 if not path.exists(cache_path):
     os.makedirs(cache_path, exist_ok=True)
 
+print("Loading FLUX pipeline...")
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", 
     torch_dtype=torch.bfloat16
 )
+
+print("Loading LoRA weights...")
 pipe.load_lora_weights(
     hf_hub_download(
         "ByteDance/Hyper-SD", 
@@ -92,9 +148,15 @@ pipe.load_lora_weights(
 )
 pipe.fuse_lora(lora_scale=0.125)
 pipe.to(device="cuda", dtype=torch.bfloat16)
-pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained(
-    "CompVis/stable-diffusion-safety-checker"
-)
+
+# Safety checker initialization
+try:
+    pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker"
+    )
+except Exception as e:
+    print(f"Warning: Could not load safety checker: {e}")
+    pipe.safety_checker = None
 
 @spaces.GPU
 def generate_caption(image, model_name='gokaygokay/Florence-2-Flux-Large'):
@@ -142,6 +204,7 @@ def process_and_save_image(height, width, steps, scales, prompt, seed):
             return generated_image
         except Exception as e:
             print(f"Error in image generation: {str(e)}")
+            gr.Warning(f"Error generating image: {str(e)}")
             return None
 
 def get_random_seed():