Update app.py

app.py

@@ -1,49 +1,13 @@
+# No need for dynamic package installation since we're using requirements.txt
 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",
-        "timm",  # Required for Florence-2
-    ]
-    
-    for package in packages:
-        try:
-            subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", package], check=True)
-        except Exception as e:
-            print(f"Warning: Could not install {package}: {e}")
-
-# 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
+# Optional: Try to install flash-attn if not present (but don't fail if it doesn't work)
 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 flash_attn
+    print("Flash attention is available")
+except ImportError:
+    print("Flash attention not available - using standard attention (this is fine)")
 
 import spaces
 import argparse
@@ -61,20 +25,7 @@ from diffusers.pipelines.stable_diffusion import safety_checker
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM
 
-# 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"
+# No special attention mechanisms needed - the model will work fine without them
 
 # Setup and initialization code
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
@@ -91,53 +42,52 @@ print("Initializing Florence models...")
 florence_models = {}
 florence_processors = {}
 
+# Check if timm is available
 try:
-    # Try importing timm to verify it's available
     import timm
-    print("timm library available")
+    print("timm library is available")
 except ImportError:
-    print("Installing timm...")
+    print("ERROR: timm library not found. Installing...")
     subprocess.run([sys.executable, "-m", "pip", "install", "timm"], check=True)
     import timm
+    print("timm installed successfully")
 
-# Initialize Florence models with error handling
-model_names = ['gokaygokay/Florence-2-Flux-Large', 'gokaygokay/Florence-2-Flux']
+# Initialize Florence models with better error handling
+model_configs = [
+    ('gokaygokay/Florence-2-Flux-Large', True),  # Primary model
+    ('gokaygokay/Florence-2-Flux', False),        # Fallback model
+]
 
-for model_name in model_names:
+for model_name, is_primary in model_configs:
     try:
         print(f"Loading {model_name}...")
         florence_models[model_name] = AutoModelForCausalLM.from_pretrained(
             model_name, 
-            trust_remote_code=True
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,  # Add memory optimization
         ).eval()
+        
         florence_processors[model_name] = AutoProcessor.from_pretrained(
             model_name, 
             trust_remote_code=True
         )
         print(f"Successfully loaded {model_name}")
+        
+        # If we successfully loaded the primary model, we can skip the fallback
+        if is_primary:
+            break
+            
     except Exception as e:
         print(f"Warning: Could not load {model_name}: {e}")
-        # If the large model fails, we'll fall back to the smaller one
-        if model_name == 'gokaygokay/Florence-2-Flux-Large' and len(florence_models) == 0:
+        if is_primary:
             print("Attempting to load fallback model...")
-            try:
-                fallback_model = 'gokaygokay/Florence-2-Flux'
-                florence_models[model_name] = AutoModelForCausalLM.from_pretrained(
-                    fallback_model, 
-                    trust_remote_code=True
-                ).eval()
-                florence_processors[model_name] = AutoProcessor.from_pretrained(
-                    fallback_model, 
-                    trust_remote_code=True
-                )
-                print(f"Using {fallback_model} as fallback")
-            except Exception as e2:
-                print(f"Error loading fallback model: {e2}")
+        continue
 
 if not florence_models:
-    print("ERROR: No Florence models could be loaded. Caption generation will not work.")
+    print("WARNING: No Florence models could be loaded. Caption generation will not be available.")
+    print("Users will need to manually enter image descriptions.")
 else:
-    print(f"Loaded {len(florence_models)} Florence model(s)")
+    print(f"Successfully loaded {len(florence_models)} Florence model(s)")
 
 def filter_prompt(prompt):
     inappropriate_keywords = [
@@ -172,18 +122,6 @@ pipe = FluxPipeline.from_pretrained(
     torch_dtype=torch.bfloat16
 )
 
-# Configure attention mechanism
-if ATTN_METHOD == "xformers":
-    try:
-        pipe.enable_xformers_memory_efficient_attention()
-        print("Enabled xformers memory efficient attention")
-    except Exception as e:
-        print(f"Could not enable xformers: {e}")
-elif ATTN_METHOD == "flash_attn":
-    print("Flash attention available")
-else:
-    print("Using standard attention")
-
 print("Loading LoRA weights...")
 pipe.load_lora_weights(
     hf_hub_download(
@@ -408,10 +346,18 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 visible=False
             )
             
-            caption_button = gr.Button(
-                "🔍 Generate Caption from Image",
-                elem_classes=["generate-btn"]
-            )
+            # Caption button - only show if models are available
+            if florence_models:
+                caption_button = gr.Button(
+                    "🔍 Generate Caption from Image",
+                    elem_classes=["generate-btn"]
+                )
+            else:
+                caption_button = gr.Button(
+                    "⚠️ Caption Generation Unavailable - Enter Description Manually",
+                    elem_classes=["generate-btn"],
+                    interactive=False
+                )
             
             # 구분선
             gr.HTML('<hr style="margin: 20px 0;">')
@@ -480,11 +426,12 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             )
 
     # Event handlers
-    caption_button.click(
-        generate_caption,
-        inputs=[input_image, florence_model],
-        outputs=[prompt]
-    )
+    if florence_models:
+        caption_button.click(
+            generate_caption,
+            inputs=[input_image, florence_model],
+            outputs=[prompt]
+        )
     
     generate_btn.click(
         process_and_save_image,