quick_fix.py

"""
Quick Setup Script for Byte Dream
Fixes the model loading issue and helps upload to Hugging Face
"""

import os
from pathlib import Path


def check_model_exists():
    """Check if trained model exists"""
    model_paths = [
        "./models/bytedream",
        "./models",
        "./bytedream",
    ]
    
    for path in model_paths:
        if Path(path).exists():
            print(f"✓ Found model at: {path}")
            return path
    
    print("⚠ No trained model found!")
    print("\nTo train the model, run:")
    print("  python train.py --epochs 1000 --batch_size 4")
    print("\nOr download pretrained weights from Hugging Face.")
    return None


def test_inference():
    """Test inference with random initialization (no model needed)"""
    print("\n" + "="*60)
    print("Testing Byte Dream with random initialization")
    print("="*60)
    
    try:
        from bytedream.generator import ByteDreamGenerator
        
        # Initialize without model path (will use random weights)
        generator = ByteDreamGenerator(
            model_path=None,  # No pretrained model
            config_path="config.yaml",
            device="cpu",
        )
        
        print("\nGenerating test image with random weights...")
        print("(This will produce random noise, but tests the pipeline)")
        
        image = generator.generate(
            prompt="A test image",
            width=256,
            height=256,
            num_inference_steps=10,  # Fast test
        )
        
        image.save("test_output.png")
        print(f"\n✓ Test image saved to: test_output.png")
        print("\nNote: This image looks like noise because we're using random weights.")
        print("To generate meaningful images, you need to train the model first.")
        
        return True
        
    except Exception as e:
        print(f"\n❌ Error during test: {e}")
        import traceback
        traceback.print_exc()
        return False


def upload_to_hf_guide():
    """Guide for uploading to Hugging Face"""
    print("\n" + "="*60)
    print("Hugging Face Upload Guide")
    print("="*60)
    
    print("""
To upload your model to Hugging Face Hub:

STEP 1: Install required packages
----------------------------------
pip install huggingface_hub

STEP 2: Login to Hugging Face
------------------------------
huggingface-cli login

Then paste your token from: https://huggingface.co/settings/tokens

STEP 3: Train your model (if not done already)
-----------------------------------------------
python train.py --epochs 1000 --batch_size 4 --output_dir ./models/bytedream

STEP 4: Upload to Hugging Face
-------------------------------
python upload_to_hf.py --repo_id "YourUsername/ByteDream" --create_space

Replace 'YourUsername' with your actual Hugging Face username.

STEP 5: Update app.py to use the uploaded model
------------------------------------------------
After uploading, modify app.py to load from Hugging Face:

```python
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("YourUsername/ByteDream")
```

TIPS:
-----
- Make sure your model directory contains the trained weights
- Use --private flag if you want to keep the model private
- The --create_space option creates files for Hugging Face Spaces deployment
- Check your repository at: https://huggingface.co/YourUsername

For more help, see:
- https://huggingface.co/docs/hub/spaces
- https://huggingface.co/docs/huggingface_hub/guides/cli
    """)


def main():
    print("\n" + "="*60)
    print("Byte Dream - Quick Setup & Troubleshooting")
    print("="*60)
    
    # Check if model exists
    model_path = check_model_exists()
    
    # Test inference
    if model_path or True:  # Always test (can work without model)
        success = test_inference()
        
        if success:
            print("\n✓ Pipeline is working!")
            print("\nNext steps:")
            print("1. Train the model: python train.py")
            print("2. Or upload to Hugging Face (see guide below)")
    
    # Show upload guide
    upload_to_hf_guide()
    
    print("\n" + "="*60)
    print("Current status:")
    print("  - app.py has been fixed to handle missing models gracefully")
    print("  - You can now run: python app.py")
    print("  - Follow the upload guide above to deploy to Hugging Face")
    print("="*60)


if __name__ == "__main__":
    main()