import gradio as gr
import spaces
import torch
from diffusers import StableDiffusionUpscalePipeline
from PIL import Image

# মডেল লোড করা হচ্ছে 
model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipeline = StableDiffusionUpscalePipeline.from_pretrained(
    model_id, torch_dtype=torch.float16
)
pipeline = pipeline.to("cuda")

@spaces.GPU
def upscale_image(image):
    if image is None:
        return None
    
    # ছবির সাইজ বড় হলে OOM এরর এড়াতে ছোট করে নেওয়া হচ্ছে
    max_size = 512
    width, height = image.size
    if width > max_size or height > max_size:
        ratio = min(max_size/width, max_size/height)
        new_size = (int(width*ratio), int(height*ratio))
        image = image.resize(new_size, Image.LANCZOS)
    
    if image.mode != "RGB":
        image = image.convert("RGB")
        
    prompt = "masterpiece, best quality, high resolution, 8k, highly detailed, photorealistic"
    with torch.no_grad():
        upscaled = pipeline(prompt=prompt, image=image).images[0]
        
    return upscaled

# Gradio ইন্টারফেস তৈরি
demo = gr.Interface(
    fn=upscale_image,
    inputs=gr.Image(type="pil", label="এখানে আপনার ছবি দিন"),
    outputs=gr.Image(type="pil", label="8K কোয়ালিটির ছবি"),
    title="🔥 ZeroGPU 8K Image Upscaler 🔥",
    description="আপনার যেকোনো ছবি আপলোড করুন এবং এটি ZeroGPU ব্যবহার করে 8K কোয়ালিটিতে কনভার্ট হয়ে যাবে। (100% Free)",
    theme=gr.themes.Soft(),
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()