Fix: Use functions instead of classes for services

app.py

@@ -13,36 +13,27 @@ import sys
 # Add backend to path
 sys.path.insert(0, str(Path(__file__).parent))
 
-from backend.services.diffusion_pipeline import DiffusionPipelineManager
-from backend.services.upscaler import ESRGANUpscaler
-from backend.services.postprocess import PostProcessor
+from backend.services.diffusion_pipeline import get_pipeline_manager, enhance_image as enhance_with_diffusion
+from backend.services.upscaler import upscale_image
+from backend.services.postprocess import postprocess_image
 from backend.config import settings
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize pipelines (lazy loading)
-diffusion_pipeline = None
-upscaler = None
-post_processor = None
+# Initialize pipeline manager (lazy loading)
+pipeline_manager = None
 
 def initialize_models():
     """Initialize all models on first use"""
-    global diffusion_pipeline, upscaler, post_processor
+    global pipeline_manager
     
-    if diffusion_pipeline is None:
+    if pipeline_manager is None:
         logger.info("Loading Stable Diffusion pipeline...")
-        diffusion_pipeline = DiffusionPipelineManager()
-        
-    if upscaler is None:
-        logger.info("Loading upscaler...")
-        upscaler = ESRGANUpscaler()
-        
-    if post_processor is None:
-        post_processor = PostProcessor()
+        pipeline_manager = get_pipeline_manager()
 
-def enhance_image(
+def enhance_image_gradio(
     input_image: Image.Image,
     strength: float = 0.3,
     guidance_scale: float = 5.5,
@@ -78,7 +69,7 @@ def enhance_image(
         
         progress(0.2, desc="Enhancing with AI (this may take a few minutes)...")
         
-        enhanced = diffusion_pipeline.run_inference(
+        enhanced = pipeline_manager.run_inference(
             image=input_image,
             prompt=prompt,
             negative_prompt=settings.NEGATIVE_PROMPT,
@@ -92,13 +83,29 @@ def enhance_image(
         # Step 2: Optional Upscaling
         if use_upscaler:
             progress(0.75, desc="Upscaling image...")
-            enhanced = upscaler.upscale(enhanced)
+            enhanced = upscale_image(enhanced, scale=2)
             progress(0.85, desc="Upscaling complete!")
         
         # Step 3: Optional Post-processing
         if use_postprocess:
             progress(0.9, desc="Applying final touches...")
-            enhanced = post_processor.process(enhanced)
+            enhanced = postprocess_image(
+                enhanced,
+                clahe=True,
+                sharpen=1.0,
+                color_enhance=1.1
+            )
+            progress(0.85, desc="Upscaling complete!")
+        
+        # Step 3: Optional Post-processing
+        if use_postprocess:
+            progress(0.9, desc="Applying final touches...")
+            enhanced = postprocess_image(
+                enhanced,
+                clahe=True,
+                sharpen=1.0,
+                color_enhance=1.1
+            )
             progress(0.95, desc="Post-processing complete!")
         
         progress(1.0, desc="Done!")
@@ -196,7 +203,7 @@ with gr.Blocks(title="Architecture AI Enhancer", theme=gr.themes.Soft()) as demo
     
     # Connect the enhance button
     enhance_btn.click(
-        fn=enhance_image,
+        fn=enhance_image_gradio,
         inputs=[
             input_image,
             strength,
@@ -216,7 +223,7 @@ with gr.Blocks(title="Architecture AI Enhancer", theme=gr.themes.Soft()) as demo
         ],
         inputs=[input_image, strength, guidance_scale, custom_prompt, use_upscaler, use_postprocess],
         outputs=output_image,
-        fn=enhance_image,
+        fn=enhance_image_gradio,
         cache_examples=False
     )