Apply code formatting and fix compatibility issues

- Update pydantic dependencies for HF Spaces compatibility
- Apply ruff formatting across codebase
- Fix import organization and type annotations
- Ensure SDXL-only model architecture

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

app.py

@@ -20,6 +20,7 @@ if src_path not in sys.path:
 
 # Initialize config early to setup cache paths before model imports
 from core.config import config
+
 print(f"🔧 Environment: {'HF Spaces' if config.is_hf_spaces else 'Local'}")
 print(f"🔧 Device: {config.device}, dtype: {config.dtype}")
 
@@ -28,64 +29,66 @@ from ui.generate.generate_tab import GenerateTab
 from ui.identify.identify_tab import IdentifyTab
 from ui.train.train_tab import TrainTab
 
+
 class FlowerifyApp:
     """Main application class for Flowerify."""
-    
+
     def __init__(self):
         self.generate_tab = GenerateTab()
         self.identify_tab = IdentifyTab()
         self.train_tab = TrainTab()
         self.french_style_tab = FrenchStyleTab()
-    
+
     def create_interface(self) -> gr.Blocks:
         """Create the main Gradio interface."""
         with gr.Blocks(title="🌸 Flowerify - AI Flower Generator & Identifier") as demo:
             gr.Markdown("# 🌸 Flowerfy — Text → Image + Flower Identifier")
-            
+
             with gr.Tabs():
                 # Create each tab
                 generate_tab = self.generate_tab.create_ui()
                 identify_tab = self.identify_tab.create_ui()
                 train_tab = self.train_tab.create_ui()
                 french_style_tab = self.french_style_tab.create_ui()
-            
+
             # Wire cross-tab interactions
             self._setup_cross_tab_interactions()
-            
+
             # Initialize data on load
             demo.load(
                 self.train_tab._count_training_images,
-                outputs=[self.train_tab.data_status]
+                outputs=[self.train_tab.data_status],
             )
-        
+
         return demo
-    
+
     def _setup_cross_tab_interactions(self):
         """Setup interactions between tabs."""
         # Auto-send generated image to Identify tab
         self.generate_tab.output_image.change(
             self.identify_tab.set_image,
             inputs=self.generate_tab.output_image,
-            outputs=self.identify_tab.image_input
+            outputs=self.identify_tab.image_input,
         )
-    
+
     def launch(self, **kwargs):
         """Launch the application."""
         demo = self.create_interface()
         # Add share=True for HF Spaces compatibility
         if config.is_hf_spaces:
-            kwargs.setdefault('share', True)
+            kwargs.setdefault("share", True)
         return demo.queue().launch(**kwargs)
 
+
 def main():
     """Main entry point."""
     try:
         print("🌸 Starting Flowerify (SDXL models)")
         print("Loading models and initializing UI...")
-        
+
         app = FlowerifyApp()
         app.launch()
-        
+
     except KeyboardInterrupt:
         print("\n👋 Application stopped by user")
     except Exception as e:
@@ -93,5 +96,6 @@ def main():
         traceback.print_exc()
         sys.exit(1)
 
+
 if __name__ == "__main__":
-    main()
+    main()

app_original.py

@@ -1,13 +1,19 @@
-import os, torch, gradio as gr, json
-from diffusers import AutoPipelineForText2Image
-from transformers import pipeline, ConvNextImageProcessor, ConvNextForImageClassification, AutoImageProcessor, AutoModelForImageClassification
-from simple_train import simple_train
 import glob
-from pathlib import Path
-from PIL import Image
+import os
+
+import gradio as gr
 import numpy as np
+import torch
+from diffusers import AutoPipelineForText2Image
+from simple_train import simple_train
 from sklearn.cluster import KMeans
-
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+    ConvNextForImageClassification,
+    ConvNextImageProcessor,
+    pipeline,
+)
 
 MODEL_ID = os.getenv("MODEL_ID", "stabilityai/sdxl-turbo")
 
@@ -23,6 +29,7 @@ if device == "cuda":
 else:
     pipe.enable_attention_slicing()
 
+
 def generate(prompt, steps, width, height, seed):
     if seed is None or int(seed) < 0:
         generator = None
@@ -32,23 +39,49 @@ def generate(prompt, steps, width, height, seed):
     result = pipe(
         prompt=prompt,
         num_inference_steps=int(steps),
-        guidance_scale=0.0,   # SDXL-Turbo works best at 0.0
+        guidance_scale=0.0,  # SDXL-Turbo works best at 0.0
         width=int(width // 8) * 8,
         height=int(height // 8) * 8,
-        generator=generator
+        generator=generator,
     )
     return result.images[0]
 
 
-
 # ---------- Flower identification (zero-shot) ----------
 # Curated label set; edit/extend as you like
 FLOWER_LABELS = [
-    "rose", "tulip", "lily", "peony", "sunflower", "chrysanthemum", "carnation",
-    "orchid", "hydrangea", "daisy", "dahlia", "ranunculus", "anemone", "marigold",
-    "lavender", "magnolia", "gardenia", "camellia", "jasmine", "iris", "gerbera",
-    "zinnia", "hibiscus", "lotus", "poppy", "sweet pea", "freesia", "lisianthus",
-    "calla lily", "cherry blossom", "plumeria", "cosmos"
+    "rose",
+    "tulip",
+    "lily",
+    "peony",
+    "sunflower",
+    "chrysanthemum",
+    "carnation",
+    "orchid",
+    "hydrangea",
+    "daisy",
+    "dahlia",
+    "ranunculus",
+    "anemone",
+    "marigold",
+    "lavender",
+    "magnolia",
+    "gardenia",
+    "camellia",
+    "jasmine",
+    "iris",
+    "gerbera",
+    "zinnia",
+    "hibiscus",
+    "lotus",
+    "poppy",
+    "sweet pea",
+    "freesia",
+    "lisianthus",
+    "calla lily",
+    "cherry blossom",
+    "plumeria",
+    "cosmos",
 ]
 
 # Initialize classifier - will be updated when trained model is loaded
@@ -58,6 +91,7 @@ convnext_model = None
 convnext_processor = None
 current_model_path = "facebook/convnext-base-224-22k"
 
+
 def load_classifier(model_path="facebook/convnext-base-224-22k"):
     global zs_classifier, convnext_model, convnext_processor, current_model_path
     try:
@@ -70,149 +104,174 @@ def load_classifier(model_path="facebook/convnext-base-224-22k"):
             zs_classifier = pipeline(
                 task="zero-shot-image-classification",
                 model="openai/clip-vit-base-patch32",
-                device=clf_device
+                device=clf_device,
             )
             return f"✅ Loaded custom ConvNeXt model from: {model_path}"
         else:
             # Load default ConvNeXt model for feature extraction and fallback to CLIP for zero-shot
-            convnext_model = ConvNextForImageClassification.from_pretrained("facebook/convnext-base-224-22k")
-            convnext_processor = ConvNextImageProcessor.from_pretrained("facebook/convnext-base-224-22k")
+            convnext_model = ConvNextForImageClassification.from_pretrained(
+                "facebook/convnext-base-224-22k"
+            )
+            convnext_processor = ConvNextImageProcessor.from_pretrained(
+                "facebook/convnext-base-224-22k"
+            )
             zs_classifier = pipeline(
                 task="zero-shot-image-classification",
                 model="openai/clip-vit-base-patch32",
-                device=clf_device
+                device=clf_device,
             )
             current_model_path = "facebook/convnext-base-224-22k"
-            return f"✅ Loaded default ConvNeXt model: facebook/convnext-base-224-22k"
+            return "✅ Loaded default ConvNeXt model: facebook/convnext-base-224-22k"
     except Exception as e:
-        return f"❌ Error loading model: {str(e)}"
+        return f"❌ Error loading model: {e!s}"
+
 
 # Initialize with default model
 load_classifier()
 
+
 def identify_flowers(image, candidate_labels, top_k, min_score):
     if image is None:
         return [], "Please provide an image (upload or generate first)."
-    
+
     labels = candidate_labels if candidate_labels else FLOWER_LABELS
-    
+
     # Use ConvNeXt for feature extraction if we have a trained model, otherwise fallback to CLIP
-    if convnext_model is not None and os.path.exists(current_model_path) and current_model_path != "facebook/convnext-base-224-22k":
+    if (
+        convnext_model is not None
+        and os.path.exists(current_model_path)
+        and current_model_path != "facebook/convnext-base-224-22k"
+    ):
         try:
             # Use trained ConvNeXt model
             inputs = convnext_processor(images=image, return_tensors="pt")
             with torch.no_grad():
                 outputs = convnext_model(**inputs)
                 predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-            
+
             # Convert predictions to results format
             results = []
             for i, score in enumerate(predictions[0]):
                 if i < len(labels):
                     results.append({"label": labels[i], "score": float(score)})
-            
+
             # Sort by score
             results = sorted(results, key=lambda r: r["score"], reverse=True)
-        except Exception as e:
+        except Exception:
             # Fallback to CLIP zero-shot
             results = zs_classifier(
-                image,
-                candidate_labels=labels,
-                hypothesis_template="a photo of a {}"
+                image, candidate_labels=labels, hypothesis_template="a photo of a {}"
             )
     else:
         # Use CLIP zero-shot classification
         results = zs_classifier(
-            image,
-            candidate_labels=labels,
-            hypothesis_template="a photo of a {}"
+            image, candidate_labels=labels, hypothesis_template="a photo of a {}"
         )
-    
+
     # Filter and format results
     results = [r for r in results if r["score"] >= float(min_score)]
-    results = sorted(results, key=lambda r: r["score"], reverse=True)[:int(top_k)]
+    results = sorted(results, key=lambda r: r["score"], reverse=True)[: int(top_k)]
     table = [[r["label"], round(float(r["score"]), 4)] for r in results]
-    model_type = "ConvNeXt" if (convnext_model is not None and os.path.exists(current_model_path) and current_model_path != "facebook/convnext-base-224-22k") else "CLIP zero-shot"
+    model_type = (
+        "ConvNeXt"
+        if (
+            convnext_model is not None
+            and os.path.exists(current_model_path)
+            and current_model_path != "facebook/convnext-base-224-22k"
+        )
+        else "CLIP zero-shot"
+    )
     msg = f"Detected flowers using {model_type}."
     return table, msg
 
+
 # simple passthrough so the generated image appears in the Identify tab automatically
 def passthrough(img):
     return img
 
+
 # Training functions
 def get_available_models():
     models_dir = "training_data/trained_models"
     if not os.path.exists(models_dir):
         return ["facebook/convnext-base-224-22k (default)"]
-    
+
     models = ["facebook/convnext-base-224-22k (default)"]
     for item in os.listdir(models_dir):
         model_path = os.path.join(models_dir, item)
-        if os.path.isdir(model_path) and os.path.exists(os.path.join(model_path, "config.json")):
+        if os.path.isdir(model_path) and os.path.exists(
+            os.path.join(model_path, "config.json")
+        ):
             models.append(f"Custom: {item}")
     return models
 
+
 def count_training_images():
     images_dir = "training_data/images"
     if not os.path.exists(images_dir):
         return "Training directory not found"
-    
+
     total_images = 0
     flower_counts = {}
-    
+
     for flower_type in os.listdir(images_dir):
         flower_path = os.path.join(images_dir, flower_type)
         if os.path.isdir(flower_path):
-            image_files = glob.glob(os.path.join(flower_path, "*.jpg")) + \
-                         glob.glob(os.path.join(flower_path, "*.jpeg")) + \
-                         glob.glob(os.path.join(flower_path, "*.png")) + \
-                         glob.glob(os.path.join(flower_path, "*.webp"))
+            image_files = (
+                glob.glob(os.path.join(flower_path, "*.jpg"))
+                + glob.glob(os.path.join(flower_path, "*.jpeg"))
+                + glob.glob(os.path.join(flower_path, "*.png"))
+                + glob.glob(os.path.join(flower_path, "*.webp"))
+            )
             count = len(image_files)
             if count > 0:
                 flower_counts[flower_type] = count
                 total_images += count
-    
+
     if total_images == 0:
         return "No training images found. Add images to subdirectories in training_data/images/"
-    
+
     result = f"**Total images: {total_images}**\n\n"
     for flower_type, count in sorted(flower_counts.items()):
         result += f"- {flower_type}: {count} images\n"
-    
+
     return result
 
+
 def start_training(epochs=None, batch_size=None, learning_rate=None):
     try:
         # Check if training data exists
         images_dir = "training_data/images"
         if not os.path.exists(images_dir):
             return "❌ Training directory not found. Please create training_data/images/ and add your data."
-        
+
         # Count images
         total_images = 0
         for flower_type in os.listdir(images_dir):
             flower_path = os.path.join(images_dir, flower_type)
             if os.path.isdir(flower_path):
-                image_files = glob.glob(os.path.join(flower_path, "*.jpg")) + \
-                             glob.glob(os.path.join(flower_path, "*.jpeg")) + \
-                             glob.glob(os.path.join(flower_path, "*.png")) + \
-                             glob.glob(os.path.join(flower_path, "*.webp"))
+                image_files = (
+                    glob.glob(os.path.join(flower_path, "*.jpg"))
+                    + glob.glob(os.path.join(flower_path, "*.jpeg"))
+                    + glob.glob(os.path.join(flower_path, "*.png"))
+                    + glob.glob(os.path.join(flower_path, "*.webp"))
+                )
                 total_images += len(image_files)
-        
+
         if total_images < 10:
             return f"❌ Need at least 10 training images. Found {total_images}. Add more images to training_data/images/"
-        
+
         # Start training
         model_path = simple_train()
-        
+
         if model_path:
             return f"✅ Training completed! Model saved to: {model_path}"
         else:
             return "❌ Training failed. Check the console for details."
-            
+
     except Exception as e:
-        return f"❌ Training error: {str(e)}"
+        return f"❌ Training error: {e!s}"
+
 
 def load_trained_model(model_selection):
     if model_selection.startswith("Custom:"):
@@ -222,25 +281,26 @@ def load_trained_model(model_selection):
     else:
         return load_classifier("facebook/convnext-base-224-22k")
 
+
 # French-style arrangement functions
 def extract_dominant_colors(image, num_colors=5):
     """Extract dominant colors from an image using k-means clustering"""
     if image is None:
         return [], "No image provided"
-    
+
     # Convert PIL image to numpy array
     img_array = np.array(image)
-    
+
     # Reshape image to be a list of pixels
     pixels = img_array.reshape(-1, 3)
-    
+
     # Use k-means to find dominant colors
     kmeans = KMeans(n_clusters=num_colors, random_state=42, n_init=10)
     kmeans.fit(pixels)
-    
+
     # Get the colors and convert to RGB values
     colors = kmeans.cluster_centers_.astype(int)
-    
+
     # Convert to color names/descriptions
     color_names = []
     for color in colors:
@@ -268,54 +328,59 @@ def extract_dominant_colors(image, num_colors=5):
             color_names.append("orange")
         else:
             color_names.append("cream")
-    
+
     return color_names, colors
 
+
 def analyze_and_generate_french_style(image):
     """Analyze uploaded flower image and generate French-style arrangement"""
     if image is None:
         return None, "Please upload an image", ""
-    
+
     # Identify the flower type
     if zs_classifier is None:
         return None, "Model not loaded", ""
-    
+
     try:
         progress_log = "🔄 **Step 1/4:** Starting flower analysis...\n\n"
-        
+
         # Identify flower
         progress_log += "🔍 Identifying flower type using AI model...\n"
         results = zs_classifier(
-            image,
-            candidate_labels=FLOWER_LABELS,
-            hypothesis_template="a photo of a {}"
+            image, candidate_labels=FLOWER_LABELS, hypothesis_template="a photo of a {}"
         )
-        
+
         top_flower = results[0]["label"] if results else "flower"
         confidence = results[0]["score"] if results else 0
-        progress_log += f"✅ Identified: **{top_flower}** (confidence: {confidence:.2%})\n\n"
-        
+        progress_log += (
+            f"✅ Identified: **{top_flower}** (confidence: {confidence:.2%})\n\n"
+        )
+
         # Extract dominant colors
         progress_log += "🔄 **Step 2/4:** Analyzing color palette...\n\n"
         progress_log += "🎨 Extracting dominant colors from image...\n"
         color_names, color_rgb = extract_dominant_colors(image, num_colors=3)
-        
+
         # Create color description
         main_colors = color_names[:3]  # Top 3 colors
         color_desc = ", ".join(main_colors)
         progress_log += f"✅ Color palette: **{color_desc}**\n\n"
-        
+
         # Generate French-style prompt
-        progress_log += "🔄 **Step 3/4:** Creating French-style arrangement prompt...\n\n"
+        progress_log += (
+            "🔄 **Step 3/4:** Creating French-style arrangement prompt...\n\n"
+        )
         prompt = f"elegant French-style floral arrangement featuring {top_flower}s in {color_desc} colors, displayed in a clear crystal vase on a marble kitchen countertop, soft natural lighting, minimalist French country kitchen background, professional photography, sophisticated composition"
         progress_log += f"✅ Prompt created: *{prompt[:100]}...*\n\n"
-        
+
         # Generate the image
-        progress_log += "🔄 **Step 4/4:** Generating French-style arrangement image...\n\n"
+        progress_log += (
+            "🔄 **Step 4/4:** Generating French-style arrangement image...\n\n"
+        )
         progress_log += "🖼️ Using AI image generation (SDXL-Turbo)...\n"
         generated_image = generate(prompt, steps=4, width=1024, height=1024, seed=-1)
         progress_log += "✅ French-style arrangement generated successfully!\n\n"
-        
+
         # Create analysis summary
         analysis = f"""
 **🌸 Flower Analysis:**
@@ -330,14 +395,19 @@ def analyze_and_generate_french_style(image):
 **📋 Process Log:**
 {progress_log}
 """
-        
-        return generated_image, "✅ Analysis complete! French-style arrangement generated.", analysis
-        
+
+        return (
+            generated_image,
+            "✅ Analysis complete! French-style arrangement generated.",
+            analysis,
+        )
+
     except Exception as e:
-        error_log = f"❌ **Error occurred during processing:**\n\n{str(e)}\n\n"
-        if 'progress_log' in locals():
+        error_log = f"❌ **Error occurred during processing:**\n\n{e!s}\n\n"
+        if "progress_log" in locals():
             error_log += f"**Progress before error:**\n{progress_log}"
-        return None, f"❌ Error: {str(e)}", error_log
+        return None, f"❌ Error: {e!s}", error_log
+
 
 # ---------- UI ----------
 with gr.Blocks() as demo:
@@ -347,66 +417,113 @@ with gr.Blocks() as demo:
         with gr.TabItem("Generate"):
             with gr.Row():
                 with gr.Column():
-                    prompt = gr.Textbox(value="ikebana-style flower arrangement, soft natural light, minimalist", label="Prompt")
-                    steps  = gr.Slider(1, 8, value=4, step=1, label="Steps")
-                    width  = gr.Slider(512, 1536, value=1024, step=8, label="Width")
+                    prompt = gr.Textbox(
+                        value="ikebana-style flower arrangement, soft natural light, minimalist",
+                        label="Prompt",
+                    )
+                    steps = gr.Slider(1, 8, value=4, step=1, label="Steps")
+                    width = gr.Slider(512, 1536, value=1024, step=8, label="Width")
                     height = gr.Slider(512, 1536, value=1024, step=8, label="Height")
-                    seed   = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
-                    go     = gr.Button("Generate", variant="primary")
+                    seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
+                    go = gr.Button("Generate", variant="primary")
                 out = gr.Image(label="Result", type="pil")
 
         with gr.TabItem("Identify"):
             with gr.Row():
                 with gr.Column():
-                    img_in = gr.Image(label="Image (upload or auto-filled from 'Generate')", type="pil", interactive=True)
-                    labels_box = gr.CheckboxGroup(choices=FLOWER_LABELS, value=["rose","tulip","lily","peony","hydrangea","orchid","sunflower"], label="Candidate labels (edit as needed)")
+                    img_in = gr.Image(
+                        label="Image (upload or auto-filled from 'Generate')",
+                        type="pil",
+                        interactive=True,
+                    )
+                    labels_box = gr.CheckboxGroup(
+                        choices=FLOWER_LABELS,
+                        value=[
+                            "rose",
+                            "tulip",
+                            "lily",
+                            "peony",
+                            "hydrangea",
+                            "orchid",
+                            "sunflower",
+                        ],
+                        label="Candidate labels (edit as needed)",
+                    )
                     topk = gr.Slider(1, 15, value=7, step=1, label="Top-K")
-                    min_score = gr.Slider(0.0, 1.0, value=0.12, step=0.01, label="Min confidence")
+                    min_score = gr.Slider(
+                        0.0, 1.0, value=0.12, step=0.01, label="Min confidence"
+                    )
                     detect_btn = gr.Button("Identify Flowers", variant="primary")
                 with gr.Column():
-                    results_tbl = gr.Dataframe(headers=["Flower", "Confidence"], datatype=["str", "number"], interactive=False)
+                    results_tbl = gr.Dataframe(
+                        headers=["Flower", "Confidence"],
+                        datatype=["str", "number"],
+                        interactive=False,
+                    )
                     status = gr.Markdown()
 
         with gr.TabItem("Train Model"):
             gr.Markdown("## 🎯 Fine-tune the flower identification model")
-            gr.Markdown("Organize your training images in subdirectories by flower type in `training_data/images/`")
-            gr.Markdown("Example: `training_data/images/roses/`, `training_data/images/tulips/`, etc.")
-            
+            gr.Markdown(
+                "Organize your training images in subdirectories by flower type in `training_data/images/`"
+            )
+            gr.Markdown(
+                "Example: `training_data/images/roses/`, `training_data/images/tulips/`, etc."
+            )
+
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("### Training Data")
                     refresh_btn = gr.Button("🔄 Refresh Data Count", size="sm")
                     data_status = gr.Markdown()
-                    
+
                     gr.Markdown("### Training Parameters")
                     epochs = gr.Slider(1, 20, value=5, step=1, label="Training Epochs")
                     batch_size = gr.Slider(1, 16, value=8, step=1, label="Batch Size")
-                    learning_rate = gr.Number(value=1e-5, label="Learning Rate", precision=6)
-                    
+                    learning_rate = gr.Number(
+                        value=1e-5, label="Learning Rate", precision=6
+                    )
+
                     train_btn = gr.Button("🚀 Start Training", variant="primary")
-                    
+
                 with gr.Column():
                     gr.Markdown("### Model Management")
-                    model_dropdown = gr.Dropdown(choices=get_available_models(), value="facebook/convnext-base-224-22k (default)", label="Select Model")
+                    model_dropdown = gr.Dropdown(
+                        choices=get_available_models(),
+                        value="facebook/convnext-base-224-22k (default)",
+                        label="Select Model",
+                    )
                     refresh_models_btn = gr.Button("🔄 Refresh Models", size="sm")
-                    load_model_btn = gr.Button("📥 Load Selected Model", variant="secondary")
-                    
-                    model_status = gr.Markdown(f"**Current model:** {current_model_path}")
-                    
+                    load_model_btn = gr.Button(
+                        "📥 Load Selected Model", variant="secondary"
+                    )
+
+                    model_status = gr.Markdown(
+                        f"**Current model:** {current_model_path}"
+                    )
+
                     gr.Markdown("### Training Status")
                     training_output = gr.Markdown()
 
         with gr.TabItem("French Style arrangement"):
             gr.Markdown("## 🇫🇷 French-Style Flower Arrangements")
-            gr.Markdown("Upload a flower image and generate an elegant French-style arrangement with matching colors!")
-            
+            gr.Markdown(
+                "Upload a flower image and generate an elegant French-style arrangement with matching colors!"
+            )
+
             with gr.Row():
                 with gr.Column():
                     upload_img = gr.Image(label="Upload Flower Image", type="pil")
-                    analyze_btn = gr.Button("🎨 Analyze & Generate French Style", variant="primary", size="lg")
-                    
+                    analyze_btn = gr.Button(
+                        "🎨 Analyze & Generate French Style",
+                        variant="primary",
+                        size="lg",
+                    )
+
                 with gr.Column():
-                    french_result = gr.Image(label="Generated French-Style Arrangement", type="pil")
+                    french_result = gr.Image(
+                        label="Generated French-Style Arrangement", type="pil"
+                    )
                     french_status = gr.Markdown()
                     analysis_details = gr.Markdown()
 
@@ -415,28 +532,37 @@ with gr.Blocks() as demo:
     # Auto-send generated image to Identify tab
     out.change(passthrough, inputs=out, outputs=img_in)
     # Run identification
-    detect_btn.click(identify_flowers, [img_in, labels_box, topk, min_score], [results_tbl, status])
-    
+    detect_btn.click(
+        identify_flowers, [img_in, labels_box, topk, min_score], [results_tbl, status]
+    )
+
     # Training tab events
     refresh_btn.click(count_training_images, outputs=[data_status])
-    refresh_models_btn.click(lambda: gr.Dropdown(choices=get_available_models()), outputs=[model_dropdown])
-    load_model_btn.click(load_trained_model, inputs=[model_dropdown], outputs=[model_status])
-    train_btn.click(start_training, inputs=[epochs, batch_size, learning_rate], outputs=[training_output])
-    
+    refresh_models_btn.click(
+        lambda: gr.Dropdown(choices=get_available_models()), outputs=[model_dropdown]
+    )
+    load_model_btn.click(
+        load_trained_model, inputs=[model_dropdown], outputs=[model_status]
+    )
+    train_btn.click(
+        start_training,
+        inputs=[epochs, batch_size, learning_rate],
+        outputs=[training_output],
+    )
+
     # French Style tab events - update status during processing
     def update_french_status():
         return "🔄 Processing... Please wait while we analyze your flower image...", ""
-    
+
     analyze_btn.click(
-        update_french_status,
-        outputs=[french_status, analysis_details]
+        update_french_status, outputs=[french_status, analysis_details]
     ).then(
-        analyze_and_generate_french_style, 
-        inputs=[upload_img], 
-        outputs=[french_result, french_status, analysis_details]
+        analyze_and_generate_french_style,
+        inputs=[upload_img],
+        outputs=[french_result, french_status, analysis_details],
     )
-    
+
     # Initialize data count on load
     demo.load(count_training_images, outputs=[data_status])
 
-demo.queue().launch()
+demo.queue().launch()

requirements.txt

@@ -234,7 +234,7 @@ pydantic==2.10.6
     # via
     #   fastapi
     #   gradio
-pydantic-core==2.27.0
+pydantic-core==2.27.2
     # via pydantic
 pydub==0.25.1
     # via gradio

src/__init__.py

@@ -1 +1 @@
-# Flowerify application package
+# Flowerify application package

src/core/__init__.py

@@ -1 +1 @@
-# Core package
+# Core package

src/core/config.py

@@ -3,26 +3,29 @@ Configuration management for the application.
 """
 
 import os
+
 import torch
+
 from .constants import DEFAULT_MODEL_ID
 
+
 class AppConfig:
     """Application configuration singleton."""
-    
+
     def __init__(self):
         self._setup_device()
         self.model_id = DEFAULT_MODEL_ID
         # Auto-detect Hugging Face Spaces environment
         self.is_hf_spaces = os.getenv("SPACE_ID") is not None
         self._setup_cache_paths()
-        
+
     def _setup_device(self):
         """Setup device configuration for PyTorch."""
         if torch.cuda.is_available():
             self.device = "cuda"
             self.dtype = torch.float16
             self.clf_device = 0
-        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
             self.device = "mps"
             self.dtype = torch.float16
             self.clf_device = 0
@@ -30,7 +33,7 @@ class AppConfig:
             self.device = "cpu"
             self.dtype = torch.float32
             self.clf_device = -1
-    
+
     def _setup_cache_paths(self):
         """Setup cache paths based on environment."""
         if self.is_hf_spaces:
@@ -48,16 +51,17 @@ class AppConfig:
                 print(f"🏠 Using configured HF_HOME: {os.getenv('HF_HOME')}")
             else:
                 print("🏠 Using default Hugging Face cache")
-        
+
     @property
     def is_cuda_available(self):
         """Check if CUDA is available."""
         return torch.cuda.is_available()
-        
+
     @property
     def is_mps_available(self):
         """Check if Apple MPS is available."""
-        return hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()
+        return hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+
 
 # Global configuration instance
-config = AppConfig()
+config = AppConfig()

src/core/constants.py

@@ -6,7 +6,7 @@ import os
 # If using external SSD, models will be cached at /Volumes/extssd/huggingface/hub
 # This is configured via environment variables (see .env file and run.sh script)
 
-# Model configuration  
+# Model configuration
 DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "stabilityai/stable-diffusion-xl-base-1.0")
 FALLBACK_MODEL_ID = "stabilityai/sdxl-turbo"  # Lightweight fallback model
 DEFAULT_CONVNEXT_MODEL = "facebook/convnext-tiny-224"
@@ -62,4 +62,4 @@ DEFAULT_MIN_SCORE = 0.12
 DEFAULT_NUM_COLORS = 3
 
 # File extensions for image files
-SUPPORTED_IMAGE_EXTENSIONS = [".jpg", ".jpeg", ".png", ".webp"]
+SUPPORTED_IMAGE_EXTENSIONS = [".jpg", ".jpeg", ".png", ".webp"]

src/services/__init__.py

@@ -1 +1 @@
-# Services package
+# Services package

src/services/models/__init__.py

@@ -1 +1 @@
-# Models package
+# Models package

src/services/models/flower_classification.py

@@ -3,93 +3,120 @@ Flower classification service using ConvNeXt and CLIP models.
 """
 
 import os
+
 import torch
+from PIL import Image
 from transformers import (
-    pipeline, ConvNextImageProcessor, ConvNextForImageClassification,
-    AutoImageProcessor, AutoModelForImageClassification
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+    ConvNextForImageClassification,
+    ConvNextImageProcessor,
+    pipeline,
 )
-from PIL import Image
-from typing import List, Dict, Tuple, Optional
 
 try:
     from core.config import config
-    from core.constants import DEFAULT_CONVNEXT_MODEL, DEFAULT_CLIP_MODEL, FLOWER_LABELS, MODELS_DIR
+    from core.constants import (
+        DEFAULT_CLIP_MODEL,
+        DEFAULT_CONVNEXT_MODEL,
+        FLOWER_LABELS,
+        MODELS_DIR,
+    )
 except ImportError:
-    import sys
     import os
+    import sys
+
     sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
     from core.config import config
-    from core.constants import DEFAULT_CONVNEXT_MODEL, DEFAULT_CLIP_MODEL, FLOWER_LABELS, MODELS_DIR
+    from core.constants import (
+        DEFAULT_CLIP_MODEL,
+        DEFAULT_CONVNEXT_MODEL,
+        FLOWER_LABELS,
+        MODELS_DIR,
+    )
+
 
 class FlowerClassificationService:
     """Service for flower classification using ConvNeXt and CLIP models."""
-    
+
     def __init__(self):
         self.zs_classifier = None
         self.convnext_model = None
         self.convnext_processor = None
         self.current_model_path = DEFAULT_CONVNEXT_MODEL
         self._initialize_models()
-    
+
     def _initialize_models(self):
         """Initialize the classification models."""
         self.load_classifier()
-    
+
     def load_classifier(self, model_path: str = DEFAULT_CONVNEXT_MODEL) -> str:
         """Load classification model from path."""
         try:
             if os.path.exists(model_path):
                 # Load custom trained model
-                self.convnext_model = AutoModelForImageClassification.from_pretrained(model_path)
+                self.convnext_model = AutoModelForImageClassification.from_pretrained(
+                    model_path
+                )
                 self.convnext_processor = AutoImageProcessor.from_pretrained(model_path)
                 self.current_model_path = model_path
                 # Also keep zero-shot classifier for fallback
                 self.zs_classifier = pipeline(
                     task="zero-shot-image-classification",
                     model=DEFAULT_CLIP_MODEL,
-                    device=config.clf_device
+                    device=config.clf_device,
                 )
                 return f"✅ Loaded custom ConvNeXt model from: {model_path}"
             else:
                 # Load default ConvNeXt model for feature extraction and fallback to CLIP for zero-shot
-                self.convnext_model = ConvNextForImageClassification.from_pretrained(DEFAULT_CONVNEXT_MODEL)
-                self.convnext_processor = ConvNextImageProcessor.from_pretrained(DEFAULT_CONVNEXT_MODEL)
+                self.convnext_model = ConvNextForImageClassification.from_pretrained(
+                    DEFAULT_CONVNEXT_MODEL
+                )
+                self.convnext_processor = ConvNextImageProcessor.from_pretrained(
+                    DEFAULT_CONVNEXT_MODEL
+                )
                 self.zs_classifier = pipeline(
                     task="zero-shot-image-classification",
                     model=DEFAULT_CLIP_MODEL,
-                    device=config.clf_device
+                    device=config.clf_device,
                 )
                 self.current_model_path = DEFAULT_CONVNEXT_MODEL
                 return f"✅ Loaded default ConvNeXt model: {DEFAULT_CONVNEXT_MODEL}"
         except Exception as e:
-            return f"❌ Error loading model: {str(e)}"
-    
-    def identify_flowers(self, image: Optional[Image.Image], 
-                        candidate_labels: Optional[List[str]] = None,
-                        top_k: int = 7, min_score: float = 0.12) -> Tuple[List[List], str]:
+            return f"❌ Error loading model: {e!s}"
+
+    def identify_flowers(
+        self,
+        image: Image.Image | None,
+        candidate_labels: list[str] | None = None,
+        top_k: int = 7,
+        min_score: float = 0.12,
+    ) -> tuple[list[list], str]:
         """Identify flowers in an image."""
         if image is None:
             return [], "Please provide an image (upload or generate first)."
-        
+
         labels = candidate_labels if candidate_labels else FLOWER_LABELS
-        
+
         # Use ConvNeXt for feature extraction if we have a trained model
-        if (self.convnext_model is not None and 
-            os.path.exists(self.current_model_path) and 
-            self.current_model_path != DEFAULT_CONVNEXT_MODEL):
+        if (
+            self.convnext_model is not None
+            and os.path.exists(self.current_model_path)
+            and self.current_model_path != DEFAULT_CONVNEXT_MODEL
+        ):
             try:
                 # Use trained ConvNeXt model
                 inputs = self.convnext_processor(images=image, return_tensors="pt")
                 with torch.no_grad():
                     outputs = self.convnext_model(**inputs)
                     predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-                
+
                 # Convert predictions to results format
                 results = []
                 for i, score in enumerate(predictions[0]):
                     if i < len(labels):
                         results.append({"label": labels[i], "score": float(score)})
-                
+
                 # Sort by score
                 results = sorted(results, key=lambda r: r["score"], reverse=True)
                 model_type = "ConvNeXt"
@@ -101,35 +128,36 @@ class FlowerClassificationService:
             # Use CLIP zero-shot classification
             results = self._use_clip_classification(image, labels)
             model_type = "CLIP zero-shot"
-        
+
         # Filter and format results
         results = [r for r in results if r["score"] >= min_score]
         results = sorted(results, key=lambda r: r["score"], reverse=True)[:top_k]
         table = [[r["label"], round(float(r["score"]), 4)] for r in results]
         msg = f"Detected flowers using {model_type}."
         return table, msg
-    
-    def _use_clip_classification(self, image: Image.Image, labels: List[str]) -> List[Dict]:
+
+    def _use_clip_classification(
+        self, image: Image.Image, labels: list[str]
+    ) -> list[dict]:
         """Use CLIP zero-shot classification."""
         return self.zs_classifier(
-            image,
-            candidate_labels=labels,
-            hypothesis_template="a photo of a {}"
+            image, candidate_labels=labels, hypothesis_template="a photo of a {}"
         )
-    
-    def get_available_models(self) -> List[str]:
+
+    def get_available_models(self) -> list[str]:
         """Get list of available models."""
         models = [f"{DEFAULT_CONVNEXT_MODEL} (default)"]
-        
+
         if os.path.exists(MODELS_DIR):
             for item in os.listdir(MODELS_DIR):
                 model_path = os.path.join(MODELS_DIR, item)
-                if (os.path.isdir(model_path) and 
-                    os.path.exists(os.path.join(model_path, "config.json"))):
+                if os.path.isdir(model_path) and os.path.exists(
+                    os.path.join(model_path, "config.json")
+                ):
                     models.append(f"Custom: {item}")
-        
+
         return models
-    
+
     def load_trained_model(self, model_selection: str) -> str:
         """Load a specific trained model."""
         if model_selection.startswith("Custom:"):
@@ -139,5 +167,6 @@ class FlowerClassificationService:
         else:
             return self.load_classifier(DEFAULT_CONVNEXT_MODEL)
 
+
 # Global service instance
-flower_classifier = FlowerClassificationService()
+flower_classifier = FlowerClassificationService()

src/services/models/image_generation.py

@@ -1,9 +1,8 @@
 """Image generation service using SDXL models."""
 
-from typing import Optional
 
-import torch
 import numpy as np
+import torch
 from diffusers import AutoPipelineForText2Image
 from PIL import Image
 
@@ -18,6 +17,7 @@ except ImportError:
     from core.config import config
     from core.constants import DEFAULT_MODEL_ID, FALLBACK_MODEL_ID
 
+
 class ImageGenerationService:
     """Service for generating images using SDXL models."""
 
@@ -35,7 +35,7 @@ class ImageGenerationService:
         else:
             model_id = DEFAULT_MODEL_ID
             model_name = "SDXL"
-        
+
         # Try primary SDXL model
         try:
             print(f"🔄 Attempting to load {model_name}: {model_id}")
@@ -44,7 +44,7 @@ class ImageGenerationService:
             ).to(config.device)
             self.model_type = "SDXL"
             print(f"✅ {model_name} loaded successfully")
-            
+
             # Enable SDXL-specific optimizations
             if config.device == "cuda":
                 try:
@@ -53,10 +53,10 @@ class ImageGenerationService:
                     self.pipe.enable_attention_slicing()
             else:
                 self.pipe.enable_attention_slicing()
-                
+
         except Exception as e:
             print(f"⚠️  {model_name} failed to load: {e}")
-            
+
             # Try fallback to SDXL-Turbo if we're not on HF Spaces and not using it already
             if not config.is_hf_spaces and model_id != FALLBACK_MODEL_ID:
                 try:
@@ -66,7 +66,7 @@ class ImageGenerationService:
                     ).to(config.device)
                     self.model_type = "SDXL"
                     print("✅ SDXL-Turbo loaded successfully")
-                    
+
                     # Enable optimizations
                     if config.device == "cuda":
                         try:
@@ -78,17 +78,19 @@ class ImageGenerationService:
                     return
                 except Exception as turbo_error:
                     print(f"⚠️  SDXL-Turbo also failed to load: {turbo_error}")
-                    raise RuntimeError(f"All SDXL models failed to load. Last error: {turbo_error}")
+                    raise RuntimeError(
+                        f"All SDXL models failed to load. Last error: {turbo_error}"
+                    )
             else:
                 raise RuntimeError(f"SDXL model failed to load: {e}")
-    
+
     def generate(
         self,
         prompt: str,
         steps: int = 4,
         width: int = 1024,
         height: int = 1024,
-        seed: Optional[int] = None,
+        seed: int | None = None,
     ) -> Image.Image:
         """Generate an image from a text prompt."""
         if seed is None or seed < 0:
@@ -112,10 +114,10 @@ class ImageGenerationService:
 
         # Validate and clean the image before returning
         image = result.images[0]
-        
+
         # Convert to numpy array to check for invalid values
         img_array = np.array(image)
-        
+
         # Check for NaN or inf values and replace them
         if np.any(np.isnan(img_array)) or np.any(np.isinf(img_array)):
             print("⚠️  Warning: Image contains invalid values (NaN/inf), cleaning...")
@@ -123,12 +125,13 @@ class ImageGenerationService:
             # Ensure values are in valid range [0, 255]
             img_array = np.clip(img_array, 0, 255).astype(np.uint8)
             image = Image.fromarray(img_array)
-        
+
         return image
-        
+
     def get_model_info(self) -> str:
         """Get information about the currently loaded model."""
         return f"Model: {self.model_type} (Stable Diffusion XL)"
 
+
 # Global service instance
-image_generator = ImageGenerationService()
+image_generator = ImageGenerationService()

src/services/training/__init__.py

@@ -1 +1 @@
-# Training package
+# Training package

src/services/training/dataset.py

@@ -3,54 +3,59 @@ Dataset class for flower training data.
 """
 
 import os
+
 import torch
 from PIL import Image
 from torch.utils.data import Dataset
-from typing import List, Optional
 
-from utils.file_utils import get_image_files, get_flower_types_from_directory
+from utils.file_utils import get_flower_types_from_directory, get_image_files
+
 
 class FlowerDataset(Dataset):
     """Dataset for flower classification training."""
-    
-    def __init__(self, image_dir: str, processor, flower_labels: Optional[List[str]] = None):
+
+    def __init__(
+        self, image_dir: str, processor, flower_labels: list[str] | None = None
+    ):
         self.image_paths = []
         self.labels = []
         self.processor = processor
-        
+
         # Auto-detect flower types from directory structure if not provided
         if flower_labels is None:
             self.flower_labels = get_flower_types_from_directory(image_dir)
         else:
             self.flower_labels = flower_labels
-            
+
         self.label_to_id = {label: idx for idx, label in enumerate(self.flower_labels)}
-        
+
         # Load images from subdirectories (organized by flower type)
         for flower_type in os.listdir(image_dir):
             flower_path = os.path.join(image_dir, flower_type)
             if os.path.isdir(flower_path) and flower_type in self.label_to_id:
                 image_files = get_image_files(flower_path)
-                
+
                 for img_path in image_files:
                     self.image_paths.append(img_path)
                     self.labels.append(self.label_to_id[flower_type])
-        
-        print(f"Loaded {len(self.image_paths)} images from {len(set(self.labels))} flower types")
+
+        print(
+            f"Loaded {len(self.image_paths)} images from {len(set(self.labels))} flower types"
+        )
         print(f"Flower types: {self.flower_labels}")
-    
+
     def __len__(self):
         return len(self.image_paths)
-    
+
     def __getitem__(self, idx):
         image_path = self.image_paths[idx]
         image = Image.open(image_path).convert("RGB")
         label = self.labels[idx]
-        
+
         # Process image for ConvNeXt
         inputs = self.processor(images=image, return_tensors="pt")
-        
+
         return {
-            'pixel_values': inputs['pixel_values'].squeeze(),
-            'labels': torch.tensor(label, dtype=torch.long)
-        }
+            "pixel_values": inputs["pixel_values"].squeeze(),
+            "labels": torch.tensor(label, dtype=torch.long),
+        }

src/services/training/training_service.py

@@ -3,36 +3,38 @@ Training service for flower classification models.
 """
 
 import os
-from typing import Optional
 
 from core.constants import IMAGES_DIR
 from utils.file_utils import count_training_images
 
+
 class TrainingService:
     """Service for managing model training."""
-    
+
     def __init__(self):
         pass
-    
-    def start_training(self, epochs: int = 5, batch_size: int = 8, 
-                      learning_rate: float = 1e-5) -> str:
+
+    def start_training(
+        self, epochs: int = 5, batch_size: int = 8, learning_rate: float = 1e-5
+    ) -> str:
         """Start the training process."""
         try:
             # Check if training data exists
             if not os.path.exists(IMAGES_DIR):
                 return "❌ Training directory not found. Please create training_data/images/ and add your data."
-            
+
             # Count images
             total_images, _ = count_training_images()
-            
+
             if total_images < 10:
                 return f"❌ Need at least 10 training images. Found {total_images}. Add more images to training_data/images/"
-            
+
             # Import and run training (lazy import to avoid startup issues)
             try:
                 from training.simple_train import simple_train
+
                 model_path = simple_train()
-                
+
                 if model_path:
                     return f"✅ Training completed! Model saved to: {model_path}"
                 else:
@@ -41,17 +43,19 @@ class TrainingService:
                 # Fallback to old training method
                 try:
                     from simple_train import simple_train as legacy_train
+
                     model_path = legacy_train()
-                    
+
                     if model_path:
                         return f"✅ Training completed! Model saved to: {model_path}"
                     else:
                         return "❌ Training failed. Check the console for details."
                 except ImportError:
                     return "❌ Training module not found. Please ensure training scripts are available."
-                    
+
         except Exception as e:
-            return f"❌ Training error: {str(e)}"
+            return f"❌ Training error: {e!s}"
+
 
 # Global service instance
-training_service = TrainingService()
+training_service = TrainingService()

src/training/__init__.py

@@ -1 +1 @@
-# Training package
+# Training package

src/training/simple_train.py

@@ -3,114 +3,125 @@ Simple ConvNeXt training script without using the Transformers Trainer class.
 Refactored version of the original simple_train.py
 """
 
+import json
 import os
+
 import torch
-import torch.nn as nn
 from torch.utils.data import DataLoader
-from transformers import ConvNextImageProcessor, ConvNextForImageClassification
-import json
+from transformers import ConvNextForImageClassification, ConvNextImageProcessor
 
-from ..services.training.dataset import FlowerDataset
 from ..core.config import config
 from ..core.constants import DEFAULT_CONVNEXT_MODEL, MODELS_DIR
+from ..services.training.dataset import FlowerDataset
+
 
 def simple_train():
     """Simple ConvNeXt training function."""
     print("🌸 Simple ConvNeXt Flower Model Training")
     print("=" * 40)
-    
+
     # Check training data
     images_dir = "training_data/images"
     if not os.path.exists(images_dir):
         print("❌ Training directory not found")
         return
-    
+
     device = config.device
     print(f"Using device: {device}")
-    
+
     # Load model and processor
     model_name = DEFAULT_CONVNEXT_MODEL
     model = ConvNextForImageClassification.from_pretrained(model_name)
     processor = ConvNextImageProcessor.from_pretrained(model_name)
     model.to(device)
-    
+
     # Create dataset
     dataset = FlowerDataset(images_dir, processor)
-    
+
     if len(dataset) < 5:
         print("❌ Need at least 5 images for training")
         return
-    
+
     # Update model config for the number of classes
     if len(dataset.flower_labels) != model.config.num_labels:
         model.config.num_labels = len(dataset.flower_labels)
         # ConvNeXt uses hidden_sizes[-1] as the final hidden dimension
-        final_hidden_size = model.config.hidden_sizes[-1] if hasattr(model.config, 'hidden_sizes') else 768
-        model.classifier = torch.nn.Linear(final_hidden_size, len(dataset.flower_labels))
-    
+        final_hidden_size = (
+            model.config.hidden_sizes[-1]
+            if hasattr(model.config, "hidden_sizes")
+            else 768
+        )
+        model.classifier = torch.nn.Linear(
+            final_hidden_size, len(dataset.flower_labels)
+        )
+
     # Split dataset
     train_size = int(0.8 * len(dataset))
     train_dataset = torch.utils.data.Subset(dataset, range(train_size))
-    
+
     # Create data loader
     def simple_collate_fn(batch):
         pixel_values = []
         labels = []
-        
+
         for item in batch:
-            pixel_values.append(item['pixel_values'])
-            labels.append(item['labels'])
-        
+            pixel_values.append(item["pixel_values"])
+            labels.append(item["labels"])
+
         return {
-            'pixel_values': torch.stack(pixel_values),
-            'labels': torch.stack(labels)
+            "pixel_values": torch.stack(pixel_values),
+            "labels": torch.stack(labels),
         }
-    
-    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=simple_collate_fn)
-    
+
+    train_loader = DataLoader(
+        train_dataset, batch_size=4, shuffle=True, collate_fn=simple_collate_fn
+    )
+
     # Setup optimizer
     optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
-    
+
     # Training loop
     model.train()
     print(f"Starting training on {len(train_dataset)} samples...")
-    
+
     for epoch in range(3):
         total_loss = 0
         num_batches = 0
-        
+
         for batch_idx, batch in enumerate(train_loader):
             # Move to device
-            pixel_values = batch['pixel_values'].to(device)
-            labels = batch['labels'].to(device)
-            
+            pixel_values = batch["pixel_values"].to(device)
+            labels = batch["labels"].to(device)
+
             # Zero gradients
             optimizer.zero_grad()
-            
+
             # Forward pass
             outputs = model(pixel_values=pixel_values, labels=labels)
             loss = outputs.loss
-            
+
             # Backward pass
             loss.backward()
             optimizer.step()
-            
+
             total_loss += loss.item()
             num_batches += 1
-            
+
             if batch_idx % 2 == 0:
-                print(f"Epoch {epoch+1}, Batch {batch_idx+1}: Loss = {loss.item():.4f}")
-        
+                print(
+                    f"Epoch {epoch + 1}, Batch {batch_idx + 1}: Loss = {loss.item():.4f}"
+                )
+
         avg_loss = total_loss / num_batches if num_batches > 0 else 0
-        print(f"Epoch {epoch+1} completed. Average loss: {avg_loss:.4f}")
-    
+        print(f"Epoch {epoch + 1} completed. Average loss: {avg_loss:.4f}")
+
     # Save model
     output_dir = os.path.join(MODELS_DIR, "simple_trained_convnext")
     os.makedirs(output_dir, exist_ok=True)
-    
+
     model.save_pretrained(output_dir)
     processor.save_pretrained(output_dir)
-    
+
     # Save config
     config_data = {
         "model_name": model_name,
@@ -119,15 +130,16 @@ def simple_train():
         "batch_size": 4,
         "learning_rate": 1e-5,
         "train_samples": len(train_dataset),
-        "num_labels": len(dataset.flower_labels)
+        "num_labels": len(dataset.flower_labels),
     }
-    
+
     with open(os.path.join(output_dir, "training_config.json"), "w") as f:
         json.dump(config_data, f, indent=2)
-    
+
     print(f"✅ ConvNeXt training completed! Model saved to {output_dir}")
     return output_dir
 
+
 if __name__ == "__main__":
     try:
         simple_train()
@@ -136,4 +148,5 @@ if __name__ == "__main__":
     except Exception as e:
         print(f"❌ Training failed: {e}")
         import traceback
-        traceback.print_exc()
+
+        traceback.print_exc()

src/ui/__init__.py

@@ -1 +1 @@
-# UI package
+# UI package

src/ui/french_style/__init__.py

@@ -1 +1 @@
-# French style tab package
+# French style tab package

src/ui/french_style/french_style_tab.py

@@ -2,121 +2,129 @@
 French Style tab UI components and logic.
 """
 
+
 import gradio as gr
 from PIL import Image
-from typing import Optional, Tuple
 
 try:
+    from core.constants import DEFAULT_NUM_COLORS, FLOWER_LABELS
     from services.models.flower_classification import flower_classifier
     from services.models.image_generation import image_generator
     from utils.color_utils import extract_dominant_colors
-    from core.constants import FLOWER_LABELS, DEFAULT_NUM_COLORS
 except ImportError:
     # Handle when imported from root app.py
-    import sys
     import os
+    import sys
+
     sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+    from core.constants import DEFAULT_NUM_COLORS, FLOWER_LABELS
     from services.models.flower_classification import flower_classifier
     from services.models.image_generation import image_generator
     from utils.color_utils import extract_dominant_colors
-    from core.constants import FLOWER_LABELS, DEFAULT_NUM_COLORS
+
 
 class FrenchStyleTab:
     """UI component for the French Style tab."""
-    
+
     def __init__(self):
         pass
-    
+
     def create_ui(self) -> gr.TabItem:
         """Create the French Style tab UI."""
         with gr.TabItem("French Style arrangement") as tab:
             gr.Markdown("## 🇫🇷 French-Style Flower Arrangements")
-            gr.Markdown("Upload a flower image and generate an elegant French-style arrangement with matching colors!")
-            
+            gr.Markdown(
+                "Upload a flower image and generate an elegant French-style arrangement with matching colors!"
+            )
+
             with gr.Row():
                 with gr.Column():
                     self.upload_img = gr.Image(label="Upload Flower Image", type="pil")
                     self.analyze_btn = gr.Button(
-                        "🎨 Analyze & Generate French Style", 
-                        variant="primary", 
-                        size="lg"
+                        "🎨 Analyze & Generate French Style",
+                        variant="primary",
+                        size="lg",
                     )
-                    
+
                 with gr.Column():
                     self.french_result = gr.Image(
-                        label="Generated French-Style Arrangement", 
-                        type="pil"
+                        label="Generated French-Style Arrangement", type="pil"
                     )
                     self.french_status = gr.Markdown()
                     self.analysis_details = gr.Markdown()
-        
+
         # Wire events
         self.analyze_btn.click(
-            self._update_status,
-            outputs=[self.french_status, self.analysis_details]
+            self._update_status, outputs=[self.french_status, self.analysis_details]
         ).then(
             self.analyze_and_generate,
             inputs=[self.upload_img],
-            outputs=[self.french_result, self.french_status, self.analysis_details]
+            outputs=[self.french_result, self.french_status, self.analysis_details],
         )
-        
+
         return tab
-    
-    def _update_status(self) -> Tuple[str, str]:
+
+    def _update_status(self) -> tuple[str, str]:
         """Update status during processing."""
         return "🔄 Processing... Please wait while we analyze your flower image...", ""
-    
-    def analyze_and_generate(self, image: Optional[Image.Image]) -> Tuple[Optional[Image.Image], str, str]:
+
+    def analyze_and_generate(
+        self, image: Image.Image | None
+    ) -> tuple[Image.Image | None, str, str]:
         """Analyze uploaded flower image and generate French-style arrangement."""
         if image is None:
             return None, "Please upload an image", ""
-        
+
         # Check if classifier is loaded
         if flower_classifier.zs_classifier is None:
             return None, "Model not loaded", ""
-        
+
         try:
             progress_log = "🔄 **Step 1/4:** Starting flower analysis...\n\n"
-            
+
             # Identify flower
             progress_log += "🔍 Identifying flower type using AI model...\n"
             results = flower_classifier._use_clip_classification(image, FLOWER_LABELS)
-            
+
             top_flower = results[0]["label"] if results else "flower"
             confidence = results[0]["score"] if results else 0
-            progress_log += f"✅ Identified: **{top_flower}** (confidence: {confidence:.2%})\n\n"
-            
+            progress_log += (
+                f"✅ Identified: **{top_flower}** (confidence: {confidence:.2%})\n\n"
+            )
+
             # Extract dominant colors
             progress_log += "🔄 **Step 2/4:** Analyzing color palette...\n\n"
             progress_log += "🎨 Extracting dominant colors from image...\n"
-            color_names, color_rgb = extract_dominant_colors(image, num_colors=DEFAULT_NUM_COLORS)
-            
+            color_names, color_rgb = extract_dominant_colors(
+                image, num_colors=DEFAULT_NUM_COLORS
+            )
+
             # Create color description
             main_colors = color_names[:3]  # Top 3 colors
             color_desc = ", ".join(main_colors)
             progress_log += f"✅ Color palette: **{color_desc}**\n\n"
-            
+
             # Generate French-style prompt
-            progress_log += "🔄 **Step 3/4:** Creating French-style arrangement prompt...\n\n"
+            progress_log += (
+                "🔄 **Step 3/4:** Creating French-style arrangement prompt...\n\n"
+            )
             prompt = (
                 f"elegant French-style floral arrangement featuring {top_flower}s in {color_desc} colors, "
                 f"displayed in a clear crystal vase on a marble kitchen countertop, soft natural lighting, "
                 f"minimalist French country kitchen background, professional photography, sophisticated composition"
             )
             progress_log += f"✅ Prompt created: *{prompt[:100]}...*\n\n"
-            
+
             # Generate the image
-            progress_log += "🔄 **Step 4/4:** Generating French-style arrangement image...\n\n"
+            progress_log += (
+                "🔄 **Step 4/4:** Generating French-style arrangement image...\n\n"
+            )
             progress_log += "🖼️ Using AI image generation (SDXL-Turbo)...\n"
             generated_image = image_generator.generate(
-                prompt=prompt, 
-                steps=4, 
-                width=1024, 
-                height=1024, 
-                seed=None
+                prompt=prompt, steps=4, width=1024, height=1024, seed=None
             )
             progress_log += "✅ French-style arrangement generated successfully!\n\n"
-            
+
             # Create analysis summary
             analysis = f"""
 **🌸 Flower Analysis:**
@@ -131,11 +139,15 @@ class FrenchStyleTab:
 **📋 Process Log:**
 {progress_log}
 """
-            
-            return generated_image, "✅ Analysis complete! French-style arrangement generated.", analysis
-            
+
+            return (
+                generated_image,
+                "✅ Analysis complete! French-style arrangement generated.",
+                analysis,
+            )
+
         except Exception as e:
-            error_log = f"❌ **Error occurred during processing:**\n\n{str(e)}\n\n"
-            if 'progress_log' in locals():
+            error_log = f"❌ **Error occurred during processing:**\n\n{e!s}\n\n"
+            if "progress_log" in locals():
                 error_log += f"**Progress before error:**\n{progress_log}"
-            return None, f"❌ Error: {str(e)}", error_log
+            return None, f"❌ Error: {e!s}", error_log

src/ui/generate/__init__.py

@@ -1 +1 @@
-# Generate tab package
+# Generate tab package

src/ui/generate/generate_tab.py

@@ -2,27 +2,29 @@
 Generate tab UI components and logic.
 """
 
+
 import gradio as gr
 from PIL import Image
-from typing import Optional
 
 try:
+    from core.constants import DEFAULT_GENERATE_STEPS, DEFAULT_HEIGHT, DEFAULT_WIDTH
     from services.models.image_generation import image_generator
-    from core.constants import DEFAULT_GENERATE_STEPS, DEFAULT_WIDTH, DEFAULT_HEIGHT
 except ImportError:
     # Handle when imported from root app.py
-    import sys
     import os
+    import sys
+
     sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+    from core.constants import DEFAULT_GENERATE_STEPS, DEFAULT_HEIGHT, DEFAULT_WIDTH
     from services.models.image_generation import image_generator
-    from core.constants import DEFAULT_GENERATE_STEPS, DEFAULT_WIDTH, DEFAULT_HEIGHT
+
 
 class GenerateTab:
     """UI component for the Generate tab."""
-    
+
     def __init__(self):
         self.output_image = None
-    
+
     def create_ui(self) -> gr.TabItem:
         """Create the Generate tab UI."""
         with gr.TabItem("Generate") as tab:
@@ -30,7 +32,7 @@ class GenerateTab:
                 with gr.Column():
                     self.prompt_input = gr.Textbox(
                         value="ikebana-style flower arrangement, soft natural light, minimalist",
-                        label="Prompt"
+                        label="Prompt",
                     )
                     self.steps_input = gr.Slider(
                         1, 8, value=DEFAULT_GENERATE_STEPS, step=1, label="Steps"
@@ -45,23 +47,27 @@ class GenerateTab:
                         value=-1, precision=0, label="Seed (-1 = random)"
                     )
                     self.generate_btn = gr.Button("Generate", variant="primary")
-                
+
                 self.output_image = gr.Image(label="Result", type="pil")
-        
+
         # Wire events
         self.generate_btn.click(
             self.generate_image,
             inputs=[
-                self.prompt_input, self.steps_input, self.width_input,
-                self.height_input, self.seed_input
+                self.prompt_input,
+                self.steps_input,
+                self.width_input,
+                self.height_input,
+                self.seed_input,
             ],
-            outputs=self.output_image
+            outputs=self.output_image,
         )
-        
+
         return tab
-    
-    def generate_image(self, prompt: str, steps: int, width: int, 
-                      height: int, seed: int) -> Optional[Image.Image]:
+
+    def generate_image(
+        self, prompt: str, steps: int, width: int, height: int, seed: int
+    ) -> Image.Image | None:
         """Generate an image from the given parameters."""
         try:
             return image_generator.generate(
@@ -69,8 +75,8 @@ class GenerateTab:
                 steps=steps,
                 width=width,
                 height=height,
-                seed=seed if seed >= 0 else None
+                seed=seed if seed >= 0 else None,
             )
         except Exception as e:
-            gr.Warning(f"Error generating image: {str(e)}")
-            return None
+            gr.Warning(f"Error generating image: {e!s}")
+            return None

src/ui/identify/__init__.py

@@ -1 +1 @@
-# Identify tab package
+# Identify tab package

src/ui/identify/identify_tab.py

@@ -2,27 +2,29 @@
 Identify tab UI components and logic.
 """
 
+
 import gradio as gr
 from PIL import Image
-from typing import List, Optional, Tuple
 
 try:
+    from core.constants import DEFAULT_MIN_SCORE, DEFAULT_TOP_K, FLOWER_LABELS
     from services.models.flower_classification import flower_classifier
-    from core.constants import FLOWER_LABELS, DEFAULT_TOP_K, DEFAULT_MIN_SCORE
 except ImportError:
     # Handle when imported from root app.py
-    import sys
     import os
+    import sys
+
     sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+    from core.constants import DEFAULT_MIN_SCORE, DEFAULT_TOP_K, FLOWER_LABELS
     from services.models.flower_classification import flower_classifier
-    from core.constants import FLOWER_LABELS, DEFAULT_TOP_K, DEFAULT_MIN_SCORE
+
 
 class IdentifyTab:
     """UI component for the Identify tab."""
-    
+
     def __init__(self):
         pass
-    
+
     def create_ui(self) -> gr.TabItem:
         """Create the Identify tab UI."""
         with gr.TabItem("Identify") as tab:
@@ -31,52 +33,70 @@ class IdentifyTab:
                     self.image_input = gr.Image(
                         label="Image (upload or auto-filled from 'Generate')",
                         type="pil",
-                        interactive=True
+                        interactive=True,
                     )
                     self.labels_input = gr.CheckboxGroup(
                         choices=FLOWER_LABELS,
-                        value=["rose", "tulip", "lily", "peony", "hydrangea", "orchid", "sunflower"],
-                        label="Candidate labels (edit as needed)"
+                        value=[
+                            "rose",
+                            "tulip",
+                            "lily",
+                            "peony",
+                            "hydrangea",
+                            "orchid",
+                            "sunflower",
+                        ],
+                        label="Candidate labels (edit as needed)",
                     )
                     self.topk_input = gr.Slider(
                         1, 15, value=DEFAULT_TOP_K, step=1, label="Top-K"
                     )
                     self.min_score_input = gr.Slider(
-                        0.0, 1.0, value=DEFAULT_MIN_SCORE, step=0.01, label="Min confidence"
+                        0.0,
+                        1.0,
+                        value=DEFAULT_MIN_SCORE,
+                        step=0.01,
+                        label="Min confidence",
                     )
                     self.detect_btn = gr.Button("Identify Flowers", variant="primary")
-                
+
                 with gr.Column():
                     self.results_table = gr.Dataframe(
                         headers=["Flower", "Confidence"],
                         datatype=["str", "number"],
-                        interactive=False
+                        interactive=False,
                     )
                     self.status_output = gr.Markdown()
-        
+
         # Wire events
         self.detect_btn.click(
             self.identify_flowers,
             inputs=[
-                self.image_input, self.labels_input, 
-                self.topk_input, self.min_score_input
+                self.image_input,
+                self.labels_input,
+                self.topk_input,
+                self.min_score_input,
             ],
-            outputs=[self.results_table, self.status_output]
+            outputs=[self.results_table, self.status_output],
         )
-        
+
         return tab
-    
-    def identify_flowers(self, image: Optional[Image.Image], 
-                        candidate_labels: List[str], top_k: int, 
-                        min_score: float) -> Tuple[List[List], str]:
+
+    def identify_flowers(
+        self,
+        image: Image.Image | None,
+        candidate_labels: list[str],
+        top_k: int,
+        min_score: float,
+    ) -> tuple[list[list], str]:
         """Identify flowers in the provided image."""
         return flower_classifier.identify_flowers(
             image=image,
             candidate_labels=candidate_labels,
             top_k=top_k,
-            min_score=min_score
+            min_score=min_score,
         )
-    
-    def set_image(self, image: Optional[Image.Image]) -> Optional[Image.Image]:
+
+    def set_image(self, image: Image.Image | None) -> Image.Image | None:
         """Set the image for identification (used by other tabs)."""
-        return image
+        return image

src/ui/train/__init__.py

@@ -1 +1 @@
-# Train tab package
+# Train tab package

src/ui/train/train_tab.py

@@ -2,8 +2,8 @@
 Train Model tab UI components and logic.
 """
 
+
 import gradio as gr
-from typing import List
 
 try:
     from services.models.flower_classification import flower_classifier
@@ -11,32 +11,38 @@ try:
     from utils.file_utils import count_training_images
 except ImportError:
     # Handle when imported from root app.py
-    import sys
     import os
+    import sys
+
     sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
     from services.models.flower_classification import flower_classifier
     from services.training.training_service import training_service
     from utils.file_utils import count_training_images
 
+
 class TrainTab:
     """UI component for the Train Model tab."""
-    
+
     def __init__(self):
         pass
-    
+
     def create_ui(self) -> gr.TabItem:
         """Create the Train Model tab UI."""
         with gr.TabItem("Train Model") as tab:
             gr.Markdown("## 🎯 Fine-tune the flower identification model")
-            gr.Markdown("Organize your training images in subdirectories by flower type in `training_data/images/`")
-            gr.Markdown("Example: `training_data/images/roses/`, `training_data/images/tulips/`, etc.")
-            
+            gr.Markdown(
+                "Organize your training images in subdirectories by flower type in `training_data/images/`"
+            )
+            gr.Markdown(
+                "Example: `training_data/images/roses/`, `training_data/images/tulips/`, etc."
+            )
+
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("### Training Data")
                     self.refresh_btn = gr.Button("🔄 Refresh Data Count", size="sm")
                     self.data_status = gr.Markdown()
-                    
+
                     gr.Markdown("### Training Parameters")
                     self.epochs_input = gr.Slider(
                         1, 20, value=5, step=1, label="Training Epochs"
@@ -47,69 +53,71 @@ class TrainTab:
                     self.learning_rate_input = gr.Number(
                         value=1e-5, label="Learning Rate", precision=6
                     )
-                    
+
                     self.train_btn = gr.Button("🚀 Start Training", variant="primary")
-                    
+
                 with gr.Column():
                     gr.Markdown("### Model Management")
                     self.model_dropdown = gr.Dropdown(
                         choices=flower_classifier.get_available_models(),
                         value=f"{flower_classifier.current_model_path} (default)",
-                        label="Select Model"
+                        label="Select Model",
                     )
                     self.refresh_models_btn = gr.Button("🔄 Refresh Models", size="sm")
-                    self.load_model_btn = gr.Button("📥 Load Selected Model", variant="secondary")
-                    
+                    self.load_model_btn = gr.Button(
+                        "📥 Load Selected Model", variant="secondary"
+                    )
+
                     self.model_status = gr.Markdown(
                         f"**Current model:** {flower_classifier.current_model_path}"
                     )
-                    
+
                     gr.Markdown("### Training Status")
                     self.training_output = gr.Markdown()
-        
+
         # Wire events
         self.refresh_btn.click(self._count_training_images, outputs=[self.data_status])
         self.refresh_models_btn.click(
             self._refresh_models, outputs=[self.model_dropdown]
         )
         self.load_model_btn.click(
-            self._load_trained_model, 
-            inputs=[self.model_dropdown], 
-            outputs=[self.model_status]
+            self._load_trained_model,
+            inputs=[self.model_dropdown],
+            outputs=[self.model_status],
         )
         self.train_btn.click(
             self._start_training,
             inputs=[self.epochs_input, self.batch_size_input, self.learning_rate_input],
-            outputs=[self.training_output]
+            outputs=[self.training_output],
         )
-        
+
         return tab
-    
+
     def _count_training_images(self) -> str:
         """Count and display training images."""
         total_images, flower_counts = count_training_images()
-        
+
         if total_images == 0:
             return "No training images found. Add images to subdirectories in training_data/images/"
-        
+
         result = f"**Total images: {total_images}**\n\n"
         for flower_type, count in sorted(flower_counts.items()):
             result += f"- {flower_type}: {count} images\n"
-        
+
         return result
-    
+
     def _refresh_models(self) -> gr.Dropdown:
         """Refresh the list of available models."""
         return gr.Dropdown(choices=flower_classifier.get_available_models())
-    
+
     def _load_trained_model(self, model_selection: str) -> str:
         """Load the selected trained model."""
         return flower_classifier.load_trained_model(model_selection)
-    
-    def _start_training(self, epochs: int, batch_size: int, learning_rate: float) -> str:
+
+    def _start_training(
+        self, epochs: int, batch_size: int, learning_rate: float
+    ) -> str:
         """Start the training process."""
         return training_service.start_training(
-            epochs=epochs,
-            batch_size=batch_size, 
-            learning_rate=learning_rate
-        )
+            epochs=epochs, batch_size=batch_size, learning_rate=learning_rate
+        )

src/utils/__init__.py

@@ -1 +1 @@
-# Utils package
+# Utils package

src/utils/color_utils.py

@@ -2,38 +2,42 @@
 Color analysis utilities.
 """
 
+
 import numpy as np
 from PIL import Image
 from sklearn.cluster import KMeans
-from typing import List, Tuple, Optional
 
-def extract_dominant_colors(image: Optional[Image.Image], num_colors: int = 5) -> Tuple[List[str], np.ndarray]:
+
+def extract_dominant_colors(
+    image: Image.Image | None, num_colors: int = 5
+) -> tuple[list[str], np.ndarray]:
     """Extract dominant colors from an image using k-means clustering."""
     if image is None:
         return [], np.array([])
-    
+
     # Convert PIL image to numpy array
     img_array = np.array(image)
-    
+
     # Reshape image to be a list of pixels
     pixels = img_array.reshape(-1, 3)
-    
+
     # Use k-means to find dominant colors
     kmeans = KMeans(n_clusters=num_colors, random_state=42, n_init=10)
     kmeans.fit(pixels)
-    
+
     # Get the colors and convert to RGB values
     colors = kmeans.cluster_centers_.astype(int)
-    
+
     # Convert to color names/descriptions
     color_names = [_rgb_to_color_name(color) for color in colors]
-    
+
     return color_names, colors
 
+
 def _rgb_to_color_name(color: np.ndarray) -> str:
     """Convert RGB values to descriptive color name."""
     r, g, b = color
-    
+
     if r > 200 and g > 200 and b > 200:
         return "white"
     elif r < 50 and g < 50 and b < 50:
@@ -56,4 +60,4 @@ def _rgb_to_color_name(color: np.ndarray) -> str:
     elif r > 150 and g > 100 and b < 100:
         return "orange"
     else:
-        return "cream"
+        return "cream"

src/utils/file_utils.py

@@ -2,19 +2,21 @@
 File and directory utilities.
 """
 
-import os
 import glob
-from typing import List, Tuple
+import os
+
 try:
-    from ..core.constants import SUPPORTED_IMAGE_EXTENSIONS, IMAGES_DIR, MODELS_DIR
+    from ..core.constants import IMAGES_DIR, MODELS_DIR, SUPPORTED_IMAGE_EXTENSIONS
 except ImportError:
     # Handle direct execution
-    import sys
     import os
+    import sys
+
     sys.path.append(os.path.dirname(os.path.dirname(__file__)))
-    from core.constants import SUPPORTED_IMAGE_EXTENSIONS, IMAGES_DIR, MODELS_DIR
+    from core.constants import IMAGES_DIR, MODELS_DIR, SUPPORTED_IMAGE_EXTENSIONS
+
 
-def get_image_files(directory: str) -> List[str]:
+def get_image_files(directory: str) -> list[str]:
     """Get all image files from a directory."""
     image_files = []
     for ext in SUPPORTED_IMAGE_EXTENSIONS:
@@ -22,11 +24,12 @@ def get_image_files(directory: str) -> List[str]:
         image_files.extend(glob.glob(pattern))
     return image_files
 
-def get_flower_types_from_directory(image_dir: str = IMAGES_DIR) -> List[str]:
+
+def get_flower_types_from_directory(image_dir: str = IMAGES_DIR) -> list[str]:
     """Auto-detect flower types from directory structure."""
     if not os.path.exists(image_dir):
         return []
-    
+
     detected_types = []
     for item in os.listdir(image_dir):
         item_path = os.path.join(image_dir, item)
@@ -34,17 +37,18 @@ def get_flower_types_from_directory(image_dir: str = IMAGES_DIR) -> List[str]:
             image_files = get_image_files(item_path)
             if image_files:  # Only add if there are images
                 detected_types.append(item)
-    
+
     return sorted(detected_types)
 
-def count_training_images() -> Tuple[int, dict]:
+
+def count_training_images() -> tuple[int, dict]:
     """Count training images by flower type."""
     if not os.path.exists(IMAGES_DIR):
         return 0, {}
-    
+
     total_images = 0
     flower_counts = {}
-    
+
     for flower_type in os.listdir(IMAGES_DIR):
         flower_path = os.path.join(IMAGES_DIR, flower_type)
         if os.path.isdir(flower_path):
@@ -53,18 +57,21 @@ def count_training_images() -> Tuple[int, dict]:
             if count > 0:
                 flower_counts[flower_type] = count
                 total_images += count
-    
+
     return total_images, flower_counts
 
-def get_available_trained_models() -> List[str]:
+
+def get_available_trained_models() -> list[str]:
     """Get list of available trained models."""
     if not os.path.exists(MODELS_DIR):
         return []
-    
+
     models = []
     for item in os.listdir(MODELS_DIR):
         model_path = os.path.join(MODELS_DIR, item)
-        if os.path.isdir(model_path) and os.path.exists(os.path.join(model_path, "config.json")):
+        if os.path.isdir(model_path) and os.path.exists(
+            os.path.join(model_path, "config.json")
+        ):
             models.append(item)
-    
-    return sorted(models)
+
+    return sorted(models)

test_external_cache.py

@@ -5,59 +5,60 @@ import os
 import sys
 from pathlib import Path
 
+
 def test_cache_configuration():
     """Test that the external cache configuration is working."""
-    
+
     print("🧪 Testing External SSD Cache Configuration")
     print("=" * 50)
-    
+
     # Check if external SSD is mounted
     external_path = Path("/Volumes/extssd")
     if not external_path.exists():
         print("❌ External SSD not found at /Volumes/extssd")
         return False
-    
+
     print("✅ External SSD is mounted")
-    
+
     # Check if HF_HOME is set correctly
     hf_home = os.environ.get("HF_HOME")
     expected_hf_home = "/Volumes/extssd/huggingface"
-    
+
     if hf_home != expected_hf_home:
-        print(f"⚠️  HF_HOME not set correctly. Expected: {expected_hf_home}, Got: {hf_home}")
+        print(
+            f"⚠️  HF_HOME not set correctly. Expected: {expected_hf_home}, Got: {hf_home}"
+        )
         print("   Set HF_HOME with: export HF_HOME=/Volumes/extssd/huggingface")
         return False
-    
+
     print(f"✅ HF_HOME correctly set to: {hf_home}")
-    
+
     # Check if cache directories exist
     hub_cache = Path(hf_home) / "hub"
     if not hub_cache.exists():
         print(f"❌ Hub cache directory not found at: {hub_cache}")
         return False
-    
+
     print(f"✅ Hub cache directory exists at: {hub_cache}")
-    
+
     # Check if models are present
     model_count = len(list(hub_cache.glob("models--*")))
     print(f"✅ Found {model_count} models in cache")
-    
+
     # Test importing Hugging Face libraries and check their cache detection
     try:
-        from huggingface_hub import HfFolder
         from transformers import AutoTokenizer
-        from diffusers import DiffusionPipeline
-        
+
         print("✅ Hugging Face libraries imported successfully")
-        
+
         # Test a small model to verify cache is working
         print("🔄 Testing cache with a small model (this may take a moment)...")
-        
+
         # This should use the external cache
         tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
-        
+
         print("✅ Successfully loaded model from cache")
-        
+
         # Check if the model files are in the expected location
         clip_path = hub_cache / "models--openai--clip-vit-base-patch32"
         if clip_path.exists():
@@ -65,13 +66,14 @@ def test_cache_configuration():
         else:
             print(f"⚠️  Model files not found at expected location: {clip_path}")
             return False
-        
+
         return True
-        
+
     except Exception as e:
         print(f"❌ Error loading model: {e}")
         return False
 
+
 def main():
     """Main test function."""
     # Load .env file if available
@@ -79,12 +81,13 @@ def main():
     if env_file.exists():
         print("📁 Loading .env file...")
         from dotenv import load_dotenv
+
         load_dotenv()
     else:
         print("⚠️  No .env file found, using system environment variables")
-    
+
     success = test_cache_configuration()
-    
+
     print("\n" + "=" * 50)
     if success:
         print("🎉 All tests passed! External SSD cache is working correctly.")
@@ -93,5 +96,6 @@ def main():
         print("❌ Some tests failed. Please check the configuration.")
         sys.exit(1)
 
+
 if __name__ == "__main__":
-    main()
+    main()

tests/__init__.py

@@ -1 +1 @@
-"""Tests package for Flowerfy application."""
+"""Tests package for Flowerfy application."""

tests/test_models.py

@@ -15,18 +15,22 @@ from PIL import Image
 sys.path.append(os.path.join(os.path.dirname(os.path.dirname(__file__)), "src"))
 
 # Import all required modules - if any fail, the script will fail immediately
-from transformers import ConvNextForImageClassification, ConvNextImageProcessor, pipeline
+from transformers import (
+    ConvNextForImageClassification,
+    ConvNextImageProcessor,
+    pipeline,
+)
 
 from core.constants import DEFAULT_CLIP_MODEL, DEFAULT_CONVNEXT_MODEL
 from services.models.flower_classification import FlowerClassificationService
-from services.models.image_generation import ImageGenerationService
 
 print("✅ All dependencies imported successfully")
 
+
 def test_convnext_model() -> bool:
     """Test ConvNeXt model loading."""
     print("1️⃣ Testing ConvNeXt model loading...")
-    
+
     try:
         print(f"Loading ConvNeXt model: {DEFAULT_CONVNEXT_MODEL}")
         model = ConvNextForImageClassification.from_pretrained(DEFAULT_CONVNEXT_MODEL)
@@ -38,19 +42,23 @@ def test_convnext_model() -> bool:
         print(f"❌ ConvNeXt model test failed: {e}")
         return False
 
+
 def test_clip_model() -> bool:
     """Test CLIP model loading."""
     print("\n2️⃣ Testing CLIP model loading...")
-    
+
     try:
         print(f"Loading CLIP model: {DEFAULT_CLIP_MODEL}")
-        classifier = pipeline('zero-shot-image-classification', model=DEFAULT_CLIP_MODEL)
+        classifier = pipeline(
+            "zero-shot-image-classification", model=DEFAULT_CLIP_MODEL
+        )
         print("✅ CLIP model loaded successfully")
         return True
     except Exception as e:
         print(f"❌ CLIP model test failed: {e}")
         return False
 
+
 def test_image_generation_models() -> bool:
     """Test image generation models (SDXL models)."""
     print("\n3️⃣ Testing image generation models...")
@@ -59,42 +67,50 @@ def test_image_generation_models() -> bool:
         # Test SDXL first (now primary)
         sdxl_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
         print(f"Testing SDXL model (primary): {sdxl_model_id}")
-        
+
         try:
             from diffusers import AutoPipelineForText2Image
-            pipe = AutoPipelineForText2Image.from_pretrained(sdxl_model_id, torch_dtype=torch.float32).to("cpu")
+
+            pipe = AutoPipelineForText2Image.from_pretrained(
+                sdxl_model_id, torch_dtype=torch.float32
+            ).to("cpu")
             print("✅ SDXL model loaded successfully")
             return True
         except Exception as sdxl_error:
             print(f"⚠️  SDXL model failed: {sdxl_error}")
-            
+
             # Test SDXL-Turbo fallback
             turbo_model_id = "stabilityai/sdxl-turbo"
             print(f"Testing SDXL-Turbo fallback: {turbo_model_id}")
-            
+
             try:
-                pipe = AutoPipelineForText2Image.from_pretrained(turbo_model_id, torch_dtype=torch.float32).to("cpu")
+                pipe = AutoPipelineForText2Image.from_pretrained(
+                    turbo_model_id, torch_dtype=torch.float32
+                ).to("cpu")
                 print("✅ SDXL-Turbo model loaded successfully as fallback")
                 return True
             except Exception as turbo_error:
                 print(f"❌ Both SDXL models failed: {turbo_error}")
                 return False
-                
+
     except Exception as e:
         print(f"❌ Image generation model test failed: {e}")
         return False
 
+
 def test_flower_classification_service() -> bool:
     """Test flower classification service."""
     print("\n4️⃣ Testing flower classification service...")
-    
+
     try:
         print("Initializing flower classification service...")
         classifier = FlowerClassificationService()
-        
+
         # Create a dummy test image (3-channel RGB)
-        test_image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
-        
+        test_image = Image.fromarray(
+            np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
+        )
+
         # Test classification
         results, message = classifier.identify_flowers(test_image, top_k=3)
         print(f"✅ Classification service working: {message}")
@@ -104,10 +120,11 @@ def test_flower_classification_service() -> bool:
         print(f"❌ Classification service test failed: {e}")
         return False
 
+
 def test_image_generation_service() -> bool:
     """Test image generation service initialization."""
     print("\n5️⃣ Testing image generation service initialization...")
-    
+
     try:
         print("Testing image generation service initialization...")
         # This will test if the service can be imported and initialized
@@ -119,11 +136,12 @@ def test_image_generation_service() -> bool:
         print(f"❌ Image generation service test failed: {e}")
         return False
 
+
 def main():
     """Run all model tests."""
     print("🧪 Testing Flowerfy models...")
     print("==============================")
-    
+
     tests = [
         ("ConvNeXt Model", test_convnext_model),
         ("CLIP Model", test_clip_model),
@@ -131,10 +149,10 @@ def main():
         ("Classification Service", test_flower_classification_service),
         ("Generation Service", test_image_generation_service),
     ]
-    
+
     passed = 0
     failed = 0
-    
+
     for test_name, test_func in tests:
         try:
             if test_func():
@@ -145,11 +163,11 @@ def main():
         except Exception as e:
             failed += 1
             print(f"❌ {test_name} test failed with exception: {e}")
-    
-    print(f"\n📊 Test Results:")
+
+    print("\n📊 Test Results:")
     print(f"✅ Passed: {passed}")
     print(f"❌ Failed: {failed}")
-    
+
     if failed == 0:
         print("\n🎉 All model tests passed successfully!")
         print("======================================")
@@ -167,6 +185,7 @@ def main():
         print(f"\n❌ {failed} test(s) failed. Please check the errors above.")
         return False
 
+
 if __name__ == "__main__":
     success = main()
-    sys.exit(0 if success else 1)
+    sys.exit(0 if success else 1)

training/advanced_trainer.py

@@ -4,26 +4,32 @@ Advanced ConvNeXt training script using Transformers Trainer.
 This provides more sophisticated training features like evaluation, checkpointing, and logging.
 """
 
+import argparse
+import json
 import os
+
 import torch
-import json
-from transformers import ConvNextImageProcessor, ConvNextForImageClassification, Trainer, TrainingArguments
 from dataset import FlowerDataset, advanced_collate_fn
-import argparse
+from transformers import (
+    ConvNextForImageClassification,
+    ConvNextImageProcessor,
+    Trainer,
+    TrainingArguments,
+)
 
 
 class ConvNeXtTrainer(Trainer):
     """Custom trainer for ConvNeXt with proper loss computation."""
-    
+
     def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
         labels = inputs.get("labels")
         outputs = model(**inputs)
-        
+
         if labels is not None:
             loss = torch.nn.functional.cross_entropy(outputs.logits, labels)
         else:
             loss = outputs.loss
-        
+
         return (loss, outputs) if return_outputs else loss
 
 
@@ -34,11 +40,11 @@ def advanced_train(
     num_epochs=5,
     batch_size=8,
     learning_rate=1e-5,
-    flower_labels=None
+    flower_labels=None,
 ):
     """
     Advanced training function using Transformers Trainer.
-    
+
     Args:
         image_dir: Directory containing training images organized by flower type
         output_dir: Directory to save the trained model
@@ -47,43 +53,55 @@ def advanced_train(
         batch_size: Training batch size
         learning_rate: Learning rate for optimization
         flower_labels: List of flower labels (auto-detected if None)
-    
+
     Returns:
         str: Path to the saved model directory, or None if training failed
     """
     print("🌸 Advanced ConvNeXt Flower Model Training")
     print("=" * 50)
-    
+
     # Check training data
     if not os.path.exists(image_dir):
         print(f"❌ Training directory not found: {image_dir}")
         return None
-    
+
     # Load model and processor
     print(f"Loading model: {model_name}")
     model = ConvNextForImageClassification.from_pretrained(model_name)
     processor = ConvNextImageProcessor.from_pretrained(model_name)
-    
+
     # Create dataset
     dataset = FlowerDataset(image_dir, processor, flower_labels)
-    
+
     if len(dataset) == 0:
-        print("❌ No training data found. Please add images to subdirectories in training_data/images/")
-        print("Example: training_data/images/roses/, training_data/images/tulips/, etc.")
+        print(
+            "❌ No training data found. Please add images to subdirectories in training_data/images/"
+        )
+        print(
+            "Example: training_data/images/roses/, training_data/images/tulips/, etc."
+        )
         return None
-    
+
     # Split dataset (80% train, 20% eval)
     train_size = int(0.8 * len(dataset))
     eval_size = len(dataset) - train_size
-    train_dataset, eval_dataset = torch.utils.data.random_split(dataset, [train_size, eval_size])
-    
+    train_dataset, eval_dataset = torch.utils.data.random_split(
+        dataset, [train_size, eval_size]
+    )
+
     # Update model config for the number of classes
     if len(dataset.flower_labels) != model.config.num_labels:
         model.config.num_labels = len(dataset.flower_labels)
         # ConvNeXt uses hidden_sizes[-1] as the final hidden dimension
-        final_hidden_size = model.config.hidden_sizes[-1] if hasattr(model.config, 'hidden_sizes') else 768
-        model.classifier = torch.nn.Linear(final_hidden_size, len(dataset.flower_labels))
-    
+        final_hidden_size = (
+            model.config.hidden_sizes[-1]
+            if hasattr(model.config, "hidden_sizes")
+            else 768
+        )
+        model.classifier = torch.nn.Linear(
+            final_hidden_size, len(dataset.flower_labels)
+        )
+
     # Training arguments
     training_args = TrainingArguments(
         output_dir=output_dir,
@@ -102,7 +120,7 @@ def advanced_train(
         dataloader_num_workers=0,  # Set to 0 to avoid multiprocessing issues
         remove_unused_columns=False,
     )
-    
+
     # Create trainer
     try:
         trainer = ConvNeXtTrainer(
@@ -116,7 +134,7 @@ def advanced_train(
     except Exception as e:
         print(f"❌ Error creating trainer: {e}")
         return None
-    
+
     # Train model
     print("Starting advanced training...")
     try:
@@ -125,14 +143,15 @@ def advanced_train(
     except Exception as e:
         print(f"❌ Training failed: {e}")
         import traceback
+
         traceback.print_exc()
         return None
-    
+
     # Save final model
     final_model_path = os.path.join(output_dir, "final_model")
     model.save_pretrained(final_model_path)
     processor.save_pretrained(final_model_path)
-    
+
     # Save training config
     config = {
         "model_name": model_name,
@@ -142,27 +161,43 @@ def advanced_train(
         "learning_rate": learning_rate,
         "train_samples": len(train_dataset),
         "eval_samples": len(eval_dataset),
-        "training_type": "advanced"
+        "training_type": "advanced",
     }
-    
+
     with open(os.path.join(final_model_path, "training_config.json"), "w") as f:
         json.dump(config, f, indent=2)
-    
+
     print(f"✅ Advanced training complete! Model saved to {final_model_path}")
     return final_model_path
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Advanced ConvNeXt training for flower classification")
-    parser.add_argument("--image_dir", default="training_data/images", help="Directory containing training images")
-    parser.add_argument("--output_dir", default="training_data/trained_models/advanced_trained", help="Output directory for trained model")
-    parser.add_argument("--model_name", default="facebook/convnext-base-224-22k", help="Base model name")
-    parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs")
+    parser = argparse.ArgumentParser(
+        description="Advanced ConvNeXt training for flower classification"
+    )
+    parser.add_argument(
+        "--image_dir",
+        default="training_data/images",
+        help="Directory containing training images",
+    )
+    parser.add_argument(
+        "--output_dir",
+        default="training_data/trained_models/advanced_trained",
+        help="Output directory for trained model",
+    )
+    parser.add_argument(
+        "--model_name", default="facebook/convnext-base-224-22k", help="Base model name"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=5, help="Number of training epochs"
+    )
     parser.add_argument("--batch_size", type=int, default=8, help="Training batch size")
-    parser.add_argument("--learning_rate", type=float, default=1e-5, help="Learning rate")
-    
+    parser.add_argument(
+        "--learning_rate", type=float, default=1e-5, help="Learning rate"
+    )
+
     args = parser.parse_args()
-    
+
     try:
         result = advanced_train(
             image_dir=args.image_dir,
@@ -170,7 +205,7 @@ if __name__ == "__main__":
             model_name=args.model_name,
             num_epochs=args.epochs,
             batch_size=args.batch_size,
-            learning_rate=args.learning_rate
+            learning_rate=args.learning_rate,
         )
         if not result:
             print("❌ Training failed!")
@@ -180,5 +215,6 @@ if __name__ == "__main__":
     except Exception as e:
         print(f"❌ Training failed: {e}")
         import traceback
+
         traceback.print_exc()
-        exit(1)
+        exit(1)

training/dataset.py

@@ -3,9 +3,10 @@
 Flower Dataset class for training ConvNeXt models.
 """
 
+import glob
 import os
+
 import torch
-import glob
 from PIL import Image
 from torch.utils.data import Dataset
 
@@ -15,7 +16,7 @@ class FlowerDataset(Dataset):
         self.image_paths = []
         self.labels = []
         self.processor = processor
-        
+
         # Auto-detect flower types from directory structure if not provided
         if flower_labels is None:
             detected_types = []
@@ -28,22 +29,24 @@ class FlowerDataset(Dataset):
             self.flower_labels = sorted(detected_types)
         else:
             self.flower_labels = flower_labels
-            
+
         self.label_to_id = {label: idx for idx, label in enumerate(self.flower_labels)}
-        
+
         # Load images from subdirectories (organized by flower type)
         for flower_type in os.listdir(image_dir):
             flower_path = os.path.join(image_dir, flower_type)
             if os.path.isdir(flower_path) and flower_type in self.label_to_id:
                 image_files = self._get_image_files(flower_path)
-                
+
                 for img_path in image_files:
                     self.image_paths.append(img_path)
                     self.labels.append(self.label_to_id[flower_type])
-        
-        print(f"Loaded {len(self.image_paths)} images from {len(set(self.labels))} flower types")
+
+        print(
+            f"Loaded {len(self.image_paths)} images from {len(set(self.labels))} flower types"
+        )
         print(f"Flower types: {self.flower_labels}")
-    
+
     def _get_image_files(self, directory):
         """Get all supported image files from directory."""
         extensions = ["*.jpg", "*.jpeg", "*.png", "*.webp"]
@@ -52,21 +55,21 @@ class FlowerDataset(Dataset):
             image_files.extend(glob.glob(os.path.join(directory, ext)))
             image_files.extend(glob.glob(os.path.join(directory, ext.upper())))
         return image_files
-    
+
     def __len__(self):
         return len(self.image_paths)
-    
+
     def __getitem__(self, idx):
         image_path = self.image_paths[idx]
         image = Image.open(image_path).convert("RGB")
         label = self.labels[idx]
-        
+
         # Process image for ConvNeXt
         inputs = self.processor(images=image, return_tensors="pt")
-        
+
         return {
-            'pixel_values': inputs['pixel_values'].squeeze(),
-            'labels': torch.tensor(label, dtype=torch.long)
+            "pixel_values": inputs["pixel_values"].squeeze(),
+            "labels": torch.tensor(label, dtype=torch.long),
         }
 
 
@@ -74,29 +77,24 @@ def simple_collate_fn(batch):
     """Simple collation function for training."""
     pixel_values = []
     labels = []
-    
+
     for item in batch:
-        pixel_values.append(item['pixel_values'])
-        labels.append(item['labels'])
-    
-    return {
-        'pixel_values': torch.stack(pixel_values),
-        'labels': torch.stack(labels)
-    }
+        pixel_values.append(item["pixel_values"])
+        labels.append(item["labels"])
+
+    return {"pixel_values": torch.stack(pixel_values), "labels": torch.stack(labels)}
 
 
 def advanced_collate_fn(batch):
     """Advanced collation function for Trainer."""
     # Extract components
-    pixel_values = [item['pixel_values'] for item in batch]
-    labels = [item['labels'] for item in batch if 'labels' in item]
-    
+    pixel_values = [item["pixel_values"] for item in batch]
+    labels = [item["labels"] for item in batch if "labels" in item]
+
     # Stack everything
-    result = {
-        'pixel_values': torch.stack(pixel_values)
-    }
-    
+    result = {"pixel_values": torch.stack(pixel_values)}
+
     if labels:
-        result['labels'] = torch.stack(labels)
-    
-    return result
+        result["labels"] = torch.stack(labels)
+
+    return result

training/simple_trainer.py

@@ -4,13 +4,13 @@ Simple ConvNeXt training script without using the Transformers Trainer class.
 This is a lightweight training implementation for quick model fine-tuning.
 """
 
+import json
 import os
+
 import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-from transformers import ConvNextImageProcessor, ConvNextForImageClassification
 from dataset import FlowerDataset, simple_collate_fn
-import json
+from torch.utils.data import DataLoader
+from transformers import ConvNextForImageClassification, ConvNextImageProcessor
 
 
 def simple_train(
@@ -19,11 +19,11 @@ def simple_train(
     epochs=3,
     batch_size=4,
     learning_rate=1e-5,
-    model_name="facebook/convnext-base-224-22k"
+    model_name="facebook/convnext-base-224-22k",
 ):
     """
     Simple training function for ConvNeXt flower classification.
-    
+
     Args:
         image_dir: Directory containing training images organized by flower type
         output_dir: Directory to save the trained model
@@ -31,91 +31,107 @@ def simple_train(
         batch_size: Training batch size
         learning_rate: Learning rate for optimization
         model_name: Base ConvNeXt model to fine-tune
-    
+
     Returns:
         str: Path to the saved model directory, or None if training failed
     """
     print("🌸 Simple ConvNeXt Flower Model Training")
     print("=" * 40)
-    
+
     # Check training data
     if not os.path.exists(image_dir):
         print(f"❌ Training directory not found: {image_dir}")
         return None
-    
-    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+    device = (
+        "cuda"
+        if torch.cuda.is_available()
+        else "mps"
+        if torch.backends.mps.is_available()
+        else "cpu"
+    )
     print(f"Using device: {device}")
-    
+
     # Load model and processor
     print(f"Loading model: {model_name}")
     model = ConvNextForImageClassification.from_pretrained(model_name)
     processor = ConvNextImageProcessor.from_pretrained(model_name)
     model.to(device)
-    
+
     # Create dataset
     dataset = FlowerDataset(image_dir, processor)
-    
+
     if len(dataset) < 5:
         print("❌ Need at least 5 images for training")
         return None
-    
+
     # Split dataset
     train_size = int(0.8 * len(dataset))
     train_dataset = torch.utils.data.Subset(dataset, range(train_size))
-    
+
     # Update model config for the number of classes
     if len(dataset.flower_labels) != model.config.num_labels:
         model.config.num_labels = len(dataset.flower_labels)
         # ConvNeXt uses hidden_sizes[-1] as the final hidden dimension
-        final_hidden_size = model.config.hidden_sizes[-1] if hasattr(model.config, 'hidden_sizes') else 768
-        model.classifier = torch.nn.Linear(final_hidden_size, len(dataset.flower_labels))
+        final_hidden_size = (
+            model.config.hidden_sizes[-1]
+            if hasattr(model.config, "hidden_sizes")
+            else 768
+        )
+        model.classifier = torch.nn.Linear(
+            final_hidden_size, len(dataset.flower_labels)
+        )
         model.classifier.to(device)
-    
+
     # Create data loader
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=simple_collate_fn)
-    
+    train_loader = DataLoader(
+        train_dataset, batch_size=batch_size, shuffle=True, collate_fn=simple_collate_fn
+    )
+
     # Setup optimizer
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
-    
+
     # Training loop
     model.train()
     print(f"Starting training on {len(train_dataset)} samples for {epochs} epochs...")
-    
+
     for epoch in range(epochs):
         total_loss = 0
         num_batches = 0
-        
+
         for batch_idx, batch in enumerate(train_loader):
             # Move to device
-            pixel_values = batch['pixel_values'].to(device)
-            labels = batch['labels'].to(device)
-            
+            pixel_values = batch["pixel_values"].to(device)
+            labels = batch["labels"].to(device)
+
             # Zero gradients
             optimizer.zero_grad()
-            
+
             # Forward pass
             outputs = model(pixel_values=pixel_values, labels=labels)
             loss = outputs.loss
-            
+
             # Backward pass
             loss.backward()
             optimizer.step()
-            
+
             total_loss += loss.item()
             num_batches += 1
-            
+
             if batch_idx % 2 == 0 or batch_idx == len(train_loader) - 1:
-                print(f"Epoch {epoch+1}/{epochs}, Batch {batch_idx+1}/{len(train_loader)}: Loss = {loss.item():.4f}")
-        
+                print(
+                    f"Epoch {epoch + 1}/{epochs}, Batch {batch_idx + 1}/{len(train_loader)}: Loss = {loss.item():.4f}"
+                )
+
         avg_loss = total_loss / num_batches if num_batches > 0 else 0
-        print(f"Epoch {epoch+1} completed. Average loss: {avg_loss:.4f}")
-    
+        print(f"Epoch {epoch + 1} completed. Average loss: {avg_loss:.4f}")
+
     # Save model
     os.makedirs(output_dir, exist_ok=True)
-    
+
     model.save_pretrained(output_dir)
     processor.save_pretrained(output_dir)
-    
+
     # Save config
     config = {
         "model_name": model_name,
@@ -125,29 +141,45 @@ def simple_train(
         "learning_rate": learning_rate,
         "train_samples": len(train_dataset),
         "num_labels": len(dataset.flower_labels),
-        "training_type": "simple"
+        "training_type": "simple",
     }
-    
+
     with open(os.path.join(output_dir, "training_config.json"), "w") as f:
         json.dump(config, f, indent=2)
-    
+
     print(f"✅ ConvNeXt training completed! Model saved to {output_dir}")
     return output_dir
 
 
 if __name__ == "__main__":
     import argparse
-    
-    parser = argparse.ArgumentParser(description="Simple ConvNeXt training for flower classification")
-    parser.add_argument("--image_dir", default="training_data/images", help="Directory containing training images")
-    parser.add_argument("--output_dir", default="training_data/trained_models/simple_trained", help="Output directory for trained model")
-    parser.add_argument("--epochs", type=int, default=3, help="Number of training epochs")
+
+    parser = argparse.ArgumentParser(
+        description="Simple ConvNeXt training for flower classification"
+    )
+    parser.add_argument(
+        "--image_dir",
+        default="training_data/images",
+        help="Directory containing training images",
+    )
+    parser.add_argument(
+        "--output_dir",
+        default="training_data/trained_models/simple_trained",
+        help="Output directory for trained model",
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of training epochs"
+    )
     parser.add_argument("--batch_size", type=int, default=4, help="Training batch size")
-    parser.add_argument("--learning_rate", type=float, default=1e-5, help="Learning rate")
-    parser.add_argument("--model_name", default="facebook/convnext-base-224-22k", help="Base model name")
-    
+    parser.add_argument(
+        "--learning_rate", type=float, default=1e-5, help="Learning rate"
+    )
+    parser.add_argument(
+        "--model_name", default="facebook/convnext-base-224-22k", help="Base model name"
+    )
+
     args = parser.parse_args()
-    
+
     try:
         result = simple_train(
             image_dir=args.image_dir,
@@ -155,7 +187,7 @@ if __name__ == "__main__":
             epochs=args.epochs,
             batch_size=args.batch_size,
             learning_rate=args.learning_rate,
-            model_name=args.model_name
+            model_name=args.model_name,
         )
         if not result:
             print("❌ Training failed!")
@@ -165,5 +197,6 @@ if __name__ == "__main__":
     except Exception as e:
         print(f"❌ Training failed: {e}")
         import traceback
+
         traceback.print_exc()
-        exit(1)
+        exit(1)