import gradio as gr
from safetensors.torch import load_file
import os
import requests

# URLs y rutas de archivos
MODEL_URL = "https://huggingface.co/datasets/uruguayai/fooocus/resolve/main/juggernautXL_v8Rundiffusion.safetensors"
MODEL_PATH = "/home/user/app/juggernautXL_v8Rundiffusion.safetensors"

def download_file(url, path):
    if not os.path.exists(path):
        response = requests.get(url)
        if response.status_code == 200:
            with open(path, 'wb') as f:
                f.write(response.content)
            print(f"Downloaded {path}")
        else:
            raise Exception(f"Failed to download {url}, status code: {response.status_code}")

def load_model():
    download_file(MODEL_URL, MODEL_PATH)
    model = load_file(MODEL_PATH)  # Cargar modelo con safetensors
    return model

def predict(input_text):
    # Función de predicción simple
    return f"Model loaded and ready: {input_text}"

iface = gr.Interface(fn=predict, inputs="text", outputs="text")

if __name__ == "__main__":
    load_model()  # Asegúrate de que el modelo esté cargado
    iface.launch(server_name="0.0.0.0", server_port=7860)