File size: 4,794 Bytes
7fc1882 |
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 |
import gradio as gr
import torch
from PIL import Image
import numpy as np
from transformers import pipeline
import os
from huggingface_hub import login
import warnings
warnings.filterwarnings("ignore")
class PhotoUpscaler:
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.current_model = "microsoft/swin2SR-compressed-sr-x2-48"
self.upscaler = None
self.load_default_model()
def load_default_model(self):
"""Load default upscaling model optimized for HF Spaces"""
try:
self.upscaler = pipeline(
"image-to-image",
model=self.current_model,
device=0 if self.device == "cuda" else -1
)
return f"✅ Model loaded: {self.current_model}"
except Exception as e:
return f"❌ Error loading model: {str(e)}"
def upscale_image(self, image, scale_factor=2):
"""Upscale image with simplified processing for HF Spaces"""
if image is None:
return None, "❌ No image uploaded"
try:
# Convert to PIL if needed
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
# Limit size for HF Spaces (smaller memory limits)
max_size = 512
if max(image.size) > max_size:
ratio = max_size / max(image.size)
new_size = tuple(int(dim * ratio) for dim in image.size)
image = image.resize(new_size, Image.Resampling.LANCZOS)
# Perform upscaling
if self.upscaler:
try:
upscaled = self.upscaler(image)
if isinstance(upscaled, list):
upscaled = upscaled[0]
return upscaled, f"✅ Image upscaled using {self.current_model}"
except:
# Fallback to simple upscaling if model fails
new_size = tuple(int(dim * scale_factor) for dim in image.size)
upscaled = image.resize(new_size, Image.Resampling.LANCZOS)
return upscaled, "✅ Image upscaled using fallback method"
else:
new_size = tuple(int(dim * scale_factor) for dim in image.size)
upscaled = image.resize(new_size, Image.Resampling.LANCZOS)
return upscaled, "✅ Image upscaled using classic algorithm"
except Exception as e:
return None, f"❌ Processing error: {str(e)}"
# Initialize upscaler
upscaler = PhotoUpscaler()
def process_upscaling(image, scale_factor):
"""Main processing function for HF Spaces"""
result_image, process_msg = upscaler.upscale_image(image, scale_factor)
return result_image, process_msg
# Create Gradio interface optimized for HF Spaces
with gr.Blocks(
title="🚀 Photo Upscaler",
theme=gr.themes.Default(),
css="""
.gradio-container {
max-width: 900px !important;
margin: auto !important;
}
"""
) as demo:
gr.HTML("""
<div style="text-align: center; margin-bottom: 20px;">
<h1>🚀 Photo Upscaler</h1>
<p>Enhance your photos with AI-powered upscaling</p>
</div>
""")
with gr.Row():
with gr.Column():
gr.Markdown("### 📤 Input")
input_image = gr.Image(
label="Upload Photo",
type="pil"
)
scale_factor = gr.Slider(
minimum=1.5,
maximum=3.0,
value=2.0,
step=0.5,
label="Scale Factor"
)
upscale_btn = gr.Button(
"🔍 Upscale Image",
variant="primary"
)
with gr.Column():
gr.Markdown("### 📥 Output")
output_image = gr.Image(
label="Upscaled Image"
)
status_text = gr.Textbox(
label="Status",
interactive=False
)
# Event handler
upscale_btn.click(
fn=process_upscaling,
inputs=[input_image, scale_factor],
outputs=[output_image, status_text]
)
# Tips
gr.Markdown("""
### 💡 Tips for best results:
- Use images between 256x256 and 512x512 pixels
- PNG format provides best quality
- Lower scale factors work faster
""")
# Launch for HF Spaces
if __name__ == "__main__":
demo.launch() |