112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
"""Test Image Upscaling with Topaz API"""
|
|
import asyncio
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, '/app')
|
|
|
|
from app.database import SessionLocal
|
|
from app.models.job import Job
|
|
from app.models.asset import Asset
|
|
from app.services import image_upscaler
|
|
from uuid import uuid4
|
|
from PIL import Image
|
|
import io
|
|
|
|
async def test_upscale():
|
|
"""Test image upscaling workflow"""
|
|
db = SessionLocal()
|
|
|
|
try:
|
|
# Create a simple test image
|
|
print("📸 Creating test image...")
|
|
img = Image.new('RGB', (100, 100), color='red')
|
|
img_bytes = io.BytesIO()
|
|
img.save(img_bytes, format='PNG')
|
|
img_bytes.seek(0)
|
|
|
|
# Save test image to storage
|
|
storage_path = "/app/storage/images"
|
|
os.makedirs(storage_path, exist_ok=True)
|
|
test_filename = f"test_upscale_{uuid4()}.png"
|
|
test_filepath = os.path.join(storage_path, test_filename)
|
|
|
|
with open(test_filepath, 'wb') as f:
|
|
f.write(img_bytes.getvalue())
|
|
|
|
print(f"✓ Test image saved: {test_filepath}")
|
|
|
|
# Create input asset
|
|
input_asset = Asset(
|
|
original_filename=test_filename,
|
|
stored_filename=test_filename,
|
|
file_path=test_filepath,
|
|
file_type="image",
|
|
mime_type="image/png",
|
|
file_size_bytes=len(img_bytes.getvalue()),
|
|
width=100,
|
|
height=100,
|
|
source_module="test"
|
|
)
|
|
db.add(input_asset)
|
|
db.commit()
|
|
db.refresh(input_asset)
|
|
|
|
print(f"✓ Input asset created: {input_asset.id}")
|
|
|
|
# Create upscale job
|
|
job = Job(
|
|
module="image_upscaler",
|
|
action="upscale",
|
|
input_data={
|
|
"scale": 2,
|
|
"model": "Standard V2",
|
|
"output_format": "png",
|
|
"face_enhancement": False
|
|
},
|
|
input_asset_ids=[input_asset.id],
|
|
status="queued"
|
|
)
|
|
db.add(job)
|
|
db.commit()
|
|
db.refresh(job)
|
|
|
|
print(f"✓ Job created: {job.id}")
|
|
print(f"⏳ Starting upscale process...")
|
|
|
|
# Run upscale
|
|
await image_upscaler.upscale(str(job.id))
|
|
|
|
# Refresh job to get latest status
|
|
db.refresh(job)
|
|
|
|
print(f"\n📊 Job Status: {job.status}")
|
|
print(f"📊 Progress: {job.progress}%")
|
|
|
|
if job.status == "completed":
|
|
print(f"✅ SUCCESS!")
|
|
if job.output_asset_ids:
|
|
output_asset = db.query(Asset).filter(Asset.id == job.output_asset_ids[0]).first()
|
|
if output_asset:
|
|
print(f"✓ Output Asset ID: {output_asset.id}")
|
|
print(f"✓ Filename: {output_asset.original_filename}")
|
|
print(f"✓ Dimensions: {output_asset.width}x{output_asset.height}")
|
|
print(f"✓ File Path: {output_asset.file_path}")
|
|
print(f"✓ File Exists: {os.path.exists(output_asset.file_path)}")
|
|
else:
|
|
print("❌ Output asset not found in database")
|
|
else:
|
|
print("❌ No output_asset_ids in job")
|
|
elif job.status == "failed":
|
|
print(f"❌ FAILED: {job.error_message}")
|
|
else:
|
|
print(f"⚠️ Status: {job.status}")
|
|
|
|
except Exception as e:
|
|
print(f"❌ Test failed with error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
finally:
|
|
db.close()
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(test_upscale())
|