diff --git a/comfy_fm_newgen.py b/comfy_fm_newgen.py index 5e552af..f8d4d45 100644 --- a/comfy_fm_newgen.py +++ b/comfy_fm_newgen.py @@ -12,7 +12,12 @@ import logging.config from tqdm import tqdm from lib.rtf_parser import RTF_Parser -from lib.remove_bg import remove_bg_from_file_list +try: + from lib.remove_bg import remove_bg_from_file_list + REMBG_AVAILABLE = True +except ImportError: + REMBG_AVAILABLE = False + print("Warning: Background removal not available") from lib.generate_xml import create_config_xml, append_to_config_xml from lib.resize_images import resize_images from lib.xml_reader import extract_from_values @@ -190,9 +195,15 @@ def post_process_images( # resize_images(output_folder, processed_players) # logging.debug("Images resized successfully.") - # # Remove background from images using GPU if available - # remove_bg_from_file_list(output_folder, processed_players, use_gpu=use_gpu) - # logging.debug("Background removed from images.") + # Remove background from images if available + if REMBG_AVAILABLE: + try: + remove_bg_from_file_list(output_folder, processed_players, use_gpu=use_gpu) + logging.debug("Background removed from images.") + except Exception as e: + logging.warning(f"Background removal failed: {e}") + else: + logging.info("Background removal not available (rembg not installed). Images will have original backgrounds.") # Update or create configuration XML if update: diff --git a/gui.py b/gui.py index f489424..9b361dc 100644 --- a/gui.py +++ b/gui.py @@ -27,6 +27,14 @@ except ImportError: def parse_rtf(self, file_path): print(f"Warning: RTF parsing not available. Please install striprtf: pip install striprtf==0.0.28") return [] + +# Check if background removal is available +try: + from lib.remove_bg import remove_bg_from_file_list, REMBG_AVAILABLE + BG_REMOVAL_AVAILABLE = REMBG_AVAILABLE +except ImportError: + BG_REMOVAL_AVAILABLE = False + print("Warning: Background removal not available") from lib.generate_xml import create_config_xml, append_to_config_xml from lib.xml_reader import extract_from_values from lib.general import choose_profile, create_or_update, process_player_or_file, get_player_input @@ -400,6 +408,17 @@ class FMFaceGeneratorGUI: # Post-processing if not self.stop_generation: try: + # Remove background from images if available + if BG_REMOVAL_AVAILABLE: + try: + remove_bg_from_file_list(self.output_dir_var.get(), [item[0] for item in players_to_process], use_gpu=False) + self.log_message("Background removed from images.") + except Exception as e: + self.log_message(f"Background removal failed: {e}") + else: + self.log_message("Background removal not available (rembg not installed). Images will have original backgrounds.") + + # Update or create configuration XML if update: append_to_config_xml( self.output_dir_var.get(), diff --git a/lib/remove_bg.py b/lib/remove_bg.py index 0183dd1..20fbfc8 100644 --- a/lib/remove_bg.py +++ b/lib/remove_bg.py @@ -1,9 +1,20 @@ -from rembg import remove +try: + from rembg import remove + REMBG_AVAILABLE = True +except ImportError as e: + print(f"Warning: rembg not available: {e}") + REMBG_AVAILABLE = False + +try: + import onnxruntime as ort + ONNX_AVAILABLE = True +except ImportError: + ONNX_AVAILABLE = False + from PIL import Image from tqdm import tqdm import os from concurrent.futures import ThreadPoolExecutor -import onnxruntime as ort def process_images_in_batch(batch, directory, use_gpu): """ @@ -17,6 +28,9 @@ def process_images_in_batch(batch, directory, use_gpu): Returns: int: Number of images successfully processed in this batch. """ + if not REMBG_AVAILABLE: + return 0 + success_count = 0 for filename in batch: input_path = os.path.join(directory, f"{filename}.png") @@ -25,12 +39,17 @@ def process_images_in_batch(batch, directory, use_gpu): try: with Image.open(input_path) as img: # Initialize ONNX session options with GPU support if required - session_options = ort.SessionOptions() - providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"] - ort.set_default_logger_severity(3) # Suppress non-critical logging + if ONNX_AVAILABLE: + session_options = ort.SessionOptions() + providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"] + ort.set_default_logger_severity(3) # Suppress non-critical logging + + # Initialize the rembg remove function with appropriate providers + output = remove(img, session_options=session_options, providers=providers) + else: + # Fallback to CPU-only processing + output = remove(img) - # Initialize the rembg remove function with appropriate providers - output = remove(img, session_options=session_options, providers=providers) output.save(output_path) success_count += 1 @@ -53,6 +72,10 @@ def remove_bg_from_file_list(directory, filenames, max_workers=2, batch_size=2, Returns: int: The number of images successfully processed. """ + if not REMBG_AVAILABLE: + print("Background removal not available (rembg not installed). Skipping background removal.") + return 0 + processed_count = 0 # Divide filenames into batches diff --git a/readme.md b/readme.md index eecb55a..cf642e3 100644 --- a/readme.md +++ b/readme.md @@ -71,6 +71,10 @@ Things you will need that I will not be going over on how to setup. For the best user experience, use the graphical interface: ```bash +# Test if GUI works on your system +python test_gui.py + +# Launch the full GUI python run_gui.py ``` @@ -160,6 +164,12 @@ The tool uses Hugging Face models for image generation. You can specify which mo - If `striprtf` fails, try: `pip install striprtf==0.0.28` - Use alternative requirements: `pip install -r requirements-gui.txt` for GUI-only +**Numba/rembg compatibility issues:** +- The application now handles missing `rembg` gracefully +- Background removal is optional - images will work without it +- Test GUI functionality: `python test_gui.py` +- If you encounter Numba compilation errors, the GUI will still work + **Alternative installation methods:** - **Conda**: `conda install pytorch torchvision torchaudio cpuonly -c pytorch` - **GPU Support**: Replace `cpuonly` with `cudatoolkit=11.8` for CUDA support diff --git a/test_gui.py b/test_gui.py new file mode 100644 index 0000000..9b62c82 --- /dev/null +++ b/test_gui.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +""" +Test script to verify GUI can start without problematic dependencies +""" + +import sys +import os + +def test_imports(): + """Test if we can import the basic GUI components""" + try: + import tkinter as tk + from tkinter import ttk, filedialog, messagebox, scrolledtext + print("✓ Tkinter imports successful") + except ImportError as e: + print(f"✗ Tkinter import failed: {e}") + return False + + try: + from PIL import Image, ImageTk + print("✓ PIL imports successful") + except ImportError as e: + print(f"✗ PIL import failed: {e}") + return False + + try: + import json + import configparser + import threading + import queue + import time + print("✓ Standard library imports successful") + except ImportError as e: + print(f"✗ Standard library import failed: {e}") + return False + + # Test GUI components without problematic imports + try: + root = tk.Tk() + root.title("FM Face Generator - Test") + root.geometry("400x300") + + # Test basic widgets + ttk.Label(root, text="GUI Test Successful!").pack(pady=20) + ttk.Button(root, text="Close", command=root.quit).pack(pady=10) + + print("✓ GUI components work correctly") + root.after(1000, root.quit) # Close after 1 second + root.mainloop() + + return True + + except Exception as e: + print(f"✗ GUI test failed: {e}") + return False + +def main(): + """Main test function""" + print("Testing FM Face Generator GUI...") + print("=" * 50) + + success = test_imports() + + if success: + print("\n🎉 GUI test passed! You can run: python run_gui.py") + print("\nNote: Some features may be limited without full dependencies:") + print("- Background removal may not work") + print("- RTF parsing may not work") + print("- But the GUI interface should load correctly") + else: + print("\n❌ GUI test failed. Please check your Python installation.") + print("Make sure you have Python 3.8+ with Tkinter support.") + + return 0 if success else 1 + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file