app.py

import os
import sys
import importlib.util
from pathlib import Path
from huggingface_hub import hf_hub_download, snapshot_download, HfFileSystem
import gradio as gr

# Your private space details
PRIVATE_SPACE_REPO = "prgarg007/ai-style-transfer-dreamsofa"  # Replace with your actual private space name
HF_TOKEN = os.getenv("HF_TOKEN")

def setup_cache_directory():
    """Setup and return cache directory for private space files"""
    cache_dir = Path("private_space_cache")
    cache_dir.mkdir(exist_ok=True)
    return cache_dir

def download_private_assets(cache_dir):
    """Download necessary files from private space"""
    try:
        print("Downloading private space assets...")
        
        # Download entire repository snapshot for directory structure
        snapshot_download(
            repo_id=PRIVATE_SPACE_REPO,
            repo_type="space",
            local_dir=cache_dir,
            token=HF_TOKEN
        )
        
        print("Successfully downloaded private space assets!")
        return True
        
    except Exception as e:
        print(f"Error downloading private assets: {str(e)}")
        return False

def load_private_app():
    """Download and import the private space app"""
    try:
        print("Setting up cache directory...")
        cache_dir = setup_cache_directory()
        
        print("Downloading private space files...")
        if not download_private_assets(cache_dir):
            return None
        
        print("Loading private app code...")
        
        # Download the main app.py file from private space
        app_path = hf_hub_download(
            repo_id=PRIVATE_SPACE_REPO,
            filename="app.py",
            repo_type="space",
            token=HF_TOKEN
        )
        
        print(f"App file downloaded to: {app_path}")
        
        # Add the cache directory to Python path so imports work
        sys.path.insert(0, str(cache_dir))
        
        # Import and execute the private app
        spec_app = importlib.util.spec_from_file_location("private_app", app_path)
        private_app_module = importlib.util.module_from_spec(spec_app)
        
        # Set environment variables that might be needed
        # Copy any environment variables from current space to the imported module
        env_vars_to_copy = ["OPENAI_API_KEY", "HF_TOKEN"]
        for var in env_vars_to_copy:
            if os.getenv(var):
                os.environ[var] = os.getenv(var)
        
        # Execute the module
        spec_app.loader.exec_module(private_app_module)
        
        print("Successfully loaded private app module!")
        return private_app_module
        
    except Exception as e:
        print(f"Error loading private app: {str(e)}")
        return None

def create_public_interface():
    """Create public interface that uses private app functionality"""
    
    # Load the private app
    private_app = load_private_app()
    
    if private_app is None:
        # Fallback interface if loading fails
        with gr.Blocks(title="AI Style Transfer - Loading Error") as error_app:
            gr.Markdown("""
            # ⚠️ Service Temporarily Unavailable
            
            We're experiencing technical difficulties loading the AI processing service.
            
            **Possible causes:**
            - Network connectivity issues  
            - Authentication problems
            - Service maintenance
            
            **Solutions:**
            1. Check that HF_TOKEN environment variable is set correctly
            2. Verify the private space name is correct  
            3. Ensure you have access to the private space
            4. Try refreshing the page in a few minutes
            
            If the problem persists, please contact support.
            """)
        return error_app
    
    # If we successfully loaded the private app, create the interface using its functions
    try:
        # The private app should have a create_interface function
        if hasattr(private_app, 'create_interface'):
            print("Using create_interface from private app...")
            return private_app.create_interface()
        
        # Alternative: if it has the main components, build interface manually
        elif hasattr(private_app, 'process_images'):
            print("Building interface using private app functions...")
            
            with gr.Blocks(title="AI Style Transfer & Room Integration", theme=gr.themes.Soft()) as app:
                gr.Markdown("""
                # 🚀 AI Powered Style Transfer & Room Integration
                ## Secure Processing Portal
                
                Transform your furniture with AI-powered style transfer and room integration!
                
                **How it works:**
                1. Upload a **swatch image** (fabric/pattern you want to apply)
                2. Upload a **product image** (furniture piece to style)  
                3. Upload a **room image** (where you want to place the furniture)
                4. Select your preferred **room style**
                5. Let AI work its magic! ✨
                
                *Your code and processing logic are securely protected*
                """)
                
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### 📤 Upload Your Images")
                        
                        swatch_input = gr.Image(
                            label="🎨 Swatch/Pattern Image",
                            type="pil",
                            height=200
                        )
                        
                        product_input = gr.Image(
                            label="📦 Furniture/Product Image",
                            type="pil", 
                            height=200
                        )
                        
                        room_input = gr.Image(
                            label="🏠 Room/Background Image",
                            type="pil",
                            height=200
                        )
                        
                        room_style = gr.Dropdown(
                            choices=["modern", "vintage", "industrial", "scandinavian", "bohemian", "rustic"],
                            value="modern",
                            label="🎭 Room Style Theme"
                        )
                        
                        process_btn = gr.Button("🤖 Transform with AI", variant="primary", size="lg")
                    
                    with gr.Column():
                        gr.Markdown("### 🎯 AI Generated Results")
                        
                        styled_product_output = gr.Image(
                            label="🎨 Step 1: Styled Product",
                            height=300
                        )
                        
                        final_output = gr.Image(
                            label="🏆 Step 2: Room Integration",
                            height=300
                        )
                        
                        status_output = gr.Textbox(
                            label="🔄 Processing Status",
                            value="Ready to transform your furniture with AI...",
                            interactive=False
                        )
                
                # Connect to the private app's processing function
                process_btn.click(
                    fn=private_app.process_images,  # Use the function from private space
                    inputs=[swatch_input, product_input, room_input, room_style],
                    outputs=[styled_product_output, final_output, status_output],
                    show_progress=True
                )
                
                gr.Markdown("""
                ### 🎨 Example Use Cases:
                - **Interior Design**: Preview how different fabrics look on your furniture
                - **E-commerce**: Show products in various styles and room settings
                - **Home Renovation**: Experiment with different design aesthetics  
                - **Furniture Customization**: Visualize custom upholstery options
                
                ### ⚡ Powered by:
                - OpenAI DALL-E for image generation
                - Advanced AI for style analysis and transfer
                - Secure code execution from private repository
                
                ### 🔒 Privacy & Security:
                - Your processing code is kept private and secure
                - Images are processed securely without permanent storage
                - Code execution happens in isolated environment
                """)
            
            return app
        
        else:
            raise Exception("Private app doesn't have expected functions")
            
    except Exception as e:
        print(f"Error creating interface with private app: {str(e)}")
        
        # Final fallback
        with gr.Blocks(title="AI Style Transfer - Configuration Error") as fallback_app:
            gr.Markdown(f"""
            # ⚠️ Configuration Error
            
            Successfully connected to private space but encountered a configuration error:
            
            `{str(e)}`
            
            Please check:
            1. Private space code structure
            2. Required functions are properly defined
            3. Environment variables are set correctly
            
            Contact support if this issue persists.
            """)
        return fallback_app


"""Main function to create and launch the public interface"""
print("Starting public interface...")
print(f"Private space: {PRIVATE_SPACE_REPO}")
print(f"HF Token available: {'Yes' if HF_TOKEN else 'No'}")

# Create the public interface
app = create_public_interface()

# Launch the app
print("Launching public interface...")
app.launch()