import os
from dotenv import load_dotenv
from auto_diffusers import AutoDiffusersGenerator

load_dotenv()

def test_generation():
    api_key = os.getenv('GOOGLE_API_KEY')
    if not api_key:
        print("No API key found in .env file")
        return
    
    print("Testing auto-diffusers code generation...")
    print("API key loaded successfully")
    
    generator = AutoDiffusersGenerator(api_key)
    
    # Test with default FLUX model
    print("\nGenerating optimized code for FLUX.1-schnell...")
    
    try:
        optimized_code = generator.generate_optimized_code(
            model_name="black-forest-labs/FLUX.1-schnell",
            prompt_text="A cat holding a sign that says hello world",
            image_size=(768, 1360),
            num_inference_steps=4,
            use_manual_specs=False
        )
        
        print("\n" + "="*60)
        print("GENERATED OPTIMIZED CODE:")
        print("="*60)
        print(optimized_code)
        print("="*60)
        
    except Exception as e:
        print(f"Error generating code: {e}")

if __name__ == "__main__":
    test_generation()