Update app.py

app.py

@@ -13,10 +13,30 @@ import torch
 import random
 from transformers import pipeline
 
-os.system("pip install -e ./controlnet_aux")
+# Fix 1: Install controlnet_aux properly
+# os.system("pip install -e ./controlnet_aux")
+# Instead, try installing directly from GitHub:
+os.system("pip install git+https://github.com/lllyasviel/ControlNet-v1-1-nightly.git@main#subdirectory=annotator")
 
-from controlnet_aux import OpenposeDetector, CannyDetector
-from depth_anything_v2.dpt import DepthAnythingV2
+# Fix 2: Better error handling for the Korean translator
+def translate_to_english(text):
+    # Check if Korean characters are present
+    if any('\uAC00' <= char <= '\uD7A3' for char in text):
+        try:
+            # Try to load the translator
+            try:
+                # First, try with from_tf=True as suggested in the error message
+                translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", from_tf=True)
+            except:
+                # If that fails, try a different model
+                translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en-m2m-100")
+            
+            return translator(text, max_length=512)[0]['translation_text']
+        except Exception as e:
+            print(f"Translation error: {e}")
+            # Return original text if translation fails
+            return text
+    return text
 
 from huggingface_hub import hf_hub_download
 
@@ -26,18 +46,19 @@ login(token=hf_token)
 
 MAX_SEED = np.iinfo(np.int32).max
 
-# 번역기 설정
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
-
-def translate_to_english(text):
-    if any('\uAC00' <= char <= '\uD7A3' for char in text):
-        return translator(text, max_length=512)[0]['translation_text']
-    return text
+# Import ControlNet processors with better error handling
+try:
+    from controlnet_aux import OpenposeDetector, CannyDetector
+except ImportError:
+    print("Failed to import from controlnet_aux, trying alternate imports...")
+    try:
+        from annotator.openpose import OpenposeDetector
+        from annotator.canny import CannyDetector
+    except ImportError:
+        print("Could not import ControlNet processors. Using fallback implementations.")
+        # Define fallback implementations if needed
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
+from depth_anything_v2.dpt import DepthAnythingV2
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 model_configs = {
@@ -66,11 +87,17 @@ controlnet = FluxMultiControlNetModel([controlnet])
 pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
 pipe.to("cuda")
 
-mode_mapping = {"캐니":0, "타일":1, "깊이":2, "블러":3, "오픈포즈":4, "그레이스케일":5, "저품질": 6}
-strength_mapping = {"캐니":0.65, "타일":0.45, "깊이":0.55, "블러":0.45, "오픈포즈":0.55, "그레이스케일":0.45, "저품질": 0.4}
+# Fixed dictionary keys to use English for consistency
+mode_mapping = {"Canny":0, "Tile":1, "Depth":2, "Blur":3, "OpenPose":4, "Grayscale":5, "LowQuality": 6}
+strength_mapping = {"Canny":0.65, "Tile":0.45, "Depth":0.55, "Blur":0.45, "OpenPose":0.55, "Grayscale":0.45, "LowQuality": 0.4}
 
-canny = CannyDetector()
-open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+# Load processors with error handling
+try:
+    canny = CannyDetector()
+    open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+except Exception as e:
+    print(f"Error loading processors: {e}")
+    # Define fallback functions if needed
 
 torch.backends.cuda.matmul.allow_tf32 = True
 pipe.vae.enable_tiling()
@@ -154,47 +181,51 @@ def resize_img(input_image, max_side=768, min_side=512, size=None,
 
 @spaces.GPU()
 def infer(cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed, progress=gr.Progress(track_tqdm=True)):
+    try:    
+        control_mode_num = mode_mapping[control_mode]
+        prompt = translate_to_english(prompt)
         
-    control_mode_num = mode_mapping[control_mode]
-    prompt = translate_to_english(prompt)
-    
-    if cond_in is None:
-        if image_in is not None:
-            image_in = resize_img(load_image(image_in))
-            if control_mode == "Canny":
-                control_image = extract_canny(image_in)
-            elif control_mode == "Depth":
-                control_image = extract_depth(image_in)
-            elif control_mode == "OpenPose":
-                control_image = extract_openpose(image_in)
-            elif control_mode == "Blur":
-                control_image = apply_gaussian_blur(image_in)
-            elif control_mode == "LowQuality":
-                control_image = add_gaussian_noise(image_in)
-            elif control_mode == "Grayscale":
-                control_image = convert_to_grayscale(image_in)
-            elif control_mode == "Tile":
-                control_image = tile(image_in)
-    else:
-        control_image = resize_img(load_image(cond_in))
-
-    width, height = control_image.size
-    
-    image = pipe(
-        prompt, 
-        control_image=[control_image],
-        control_mode=[control_mode_num],
-        width=width,
-        height=height,
-        controlnet_conditioning_scale=[control_strength],
-        num_inference_steps=inference_steps, 
-        guidance_scale=guidance_scale,
-        generator=torch.manual_seed(seed),
-    ).images[0]
-
-    torch.cuda.empty_cache() 
+        if cond_in is None:
+            if image_in is not None:
+                image_in = resize_img(load_image(image_in))
+                if control_mode == "Canny":
+                    control_image = extract_canny(image_in)
+                elif control_mode == "Depth":
+                    control_image = extract_depth(image_in)
+                elif control_mode == "OpenPose":
+                    control_image = extract_openpose(image_in)
+                elif control_mode == "Blur":
+                    control_image = apply_gaussian_blur(image_in)
+                elif control_mode == "LowQuality":
+                    control_image = add_gaussian_noise(image_in)
+                elif control_mode == "Grayscale":
+                    control_image = convert_to_grayscale(image_in)
+                elif control_mode == "Tile":
+                    control_image = tile(image_in)
+        else:
+            control_image = resize_img(load_image(cond_in))
+
+        width, height = control_image.size
+        
+        image = pipe(
+            prompt, 
+            control_image=[control_image],
+            control_mode=[control_mode_num],
+            width=width,
+            height=height,
+            controlnet_conditioning_scale=[control_strength],
+            num_inference_steps=inference_steps, 
+            guidance_scale=guidance_scale,
+            generator=torch.manual_seed(seed),
+        ).images[0]
+
+        torch.cuda.empty_cache() 
+        
+        return image, control_image, gr.update(visible=True)
     
-    return image, control_image, gr.update(visible=True)
+    except Exception as e:
+        print(f"Error in inference: {e}")
+        return None, None, gr.update(visible=True)
    
 
 css = """