Update app.py

app.py

@@ -7,11 +7,15 @@ def install_packages():
         "transformers>=4.46.0",
         "diffusers>=0.31.0",
         "accelerate>=0.26.0",
-        "huggingface-hub>=0.23.0"
+        "huggingface-hub>=0.23.0",
+        "timm",  # Required for Florence-2
     ]
     
     for package in packages:
-        subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", package], check=True)
+        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:
@@ -84,27 +88,56 @@ 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', 
-        trust_remote_code=True
-    ).eval(),
-    'gokaygokay/Florence-2-Flux': AutoModelForCausalLM.from_pretrained(
-        'gokaygokay/Florence-2-Flux', 
-        trust_remote_code=True
-    ).eval(),
-}
+florence_models = {}
+florence_processors = {}
 
-florence_processors = {
-    'gokaygokay/Florence-2-Flux-Large': AutoProcessor.from_pretrained(
-        'gokaygokay/Florence-2-Flux-Large', 
-        trust_remote_code=True
-    ),
-    'gokaygokay/Florence-2-Flux': AutoProcessor.from_pretrained(
-        'gokaygokay/Florence-2-Flux', 
-        trust_remote_code=True
-    ),
-}
+try:
+    # Try importing timm to verify it's available
+    import timm
+    print("timm library available")
+except ImportError:
+    print("Installing timm...")
+    subprocess.run([sys.executable, "-m", "pip", "install", "timm"], check=True)
+    import timm
+
+# Initialize Florence models with error handling
+model_names = ['gokaygokay/Florence-2-Flux-Large', 'gokaygokay/Florence-2-Flux']
+
+for model_name in model_names:
+    try:
+        print(f"Loading {model_name}...")
+        florence_models[model_name] = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            trust_remote_code=True
+        ).eval()
+        florence_processors[model_name] = AutoProcessor.from_pretrained(
+            model_name, 
+            trust_remote_code=True
+        )
+        print(f"Successfully loaded {model_name}")
+    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:
+            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}")
+
+if not florence_models:
+    print("ERROR: No Florence models could be loaded. Caption generation will not work.")
+else:
+    print(f"Loaded {len(florence_models)} Florence model(s)")
 
 def filter_prompt(prompt):
     inappropriate_keywords = [
@@ -139,6 +172,18 @@ 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(
@@ -160,6 +205,15 @@ except Exception as e:
 
 @spaces.GPU
 def generate_caption(image, model_name='gokaygokay/Florence-2-Flux-Large'):
+    if not florence_models:
+        gr.Warning("Caption models are not loaded. Please refresh the page.")
+        return "Caption generation unavailable - please describe your image manually"
+    
+    # Use fallback model if the requested one isn't available
+    if model_name not in florence_models:
+        model_name = list(florence_models.keys())[0]
+        print(f"Using fallback model: {model_name}")
+    
     image = Image.fromarray(image)
     task_prompt = "<DESCRIPTION>"
     prompt = task_prompt + "Describe this image in great detail."
@@ -346,10 +400,11 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             )
             
             # Florence 모델 선택 - 숨김 처리
+            available_models = list(florence_models.keys()) if florence_models else []
             florence_model = gr.Dropdown(
-                choices=list(florence_models.keys()),
+                choices=available_models,
                 label="Caption Model",
-                value='gokaygokay/Florence-2-Flux-Large',
+                value=available_models[0] if available_models else None,
                 visible=False
             )