File size: 4,577 Bytes
0e3999b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """
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()
|