app.py

import os
import requests
import subprocess
import gradio as gr
from huggingface_hub import HfApi, upload_file
from datetime import datetime

# Token Hugging Face từ biến môi trường
hf_token = os.getenv("HF_TOKEN")

# URLs cần tải
app_url = "https://huggingface.co/datasets/ArrcttacsrjksX/Deffusion/resolve/main/RunModelAppp/App/sdmaster-d9b5942LatestJan182025"
model_url = "https://huggingface.co/datasets/ArrcttacsrjksX/Deffusion/resolve/main/Model/realisticVisionV60B1_v51HyperVAE.safetensors"

# Đường dẫn lưu file
app_path = "sdRundeffusiononhuggingfacemaster-ac54e00"
model_path = "realisticVisionV60B1_v51HyperVAE.safetensors"

# Hàm tải file từ Hugging Face
def download_file(url, output_path, token):
    headers = {"Authorization": f"Bearer {token}"}
    response = requests.get(url, headers=headers, stream=True)
    response.raise_for_status()  # Kiểm tra lỗi
    with open(output_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    print(f"Downloaded: {output_path}")

# Tải các file nếu chưa tồn tại
if not os.path.exists(app_path):
    download_file(app_url, app_path, hf_token)
    subprocess.run(["chmod", "+x", app_path])  # Thay đổi quyền thực thi
if not os.path.exists(model_path):
    download_file(model_url, model_path, hf_token)

# Hàm xử lý chạy ứng dụng và lưu kết quả
def run_command_and_save(
    prompt, mode, height, width, steps, seed, cfg_scale, strength, sampling_method,
    batch_count, schedule, clip_skip, vae_tiling, vae_on_cpu, clip_on_cpu, diffusion_fa,
    control_net_cpu, canny, verbose, init_image=None, weight_type=None
):
    try:
        # Lưu ảnh đầu vào nếu được cung cấp
        init_image_path = None
        if init_image is not None:
            init_image_path = "input_image.png"
            init_image.save(init_image_path)
        # Tạo lệnh chạy
        command = [
            f"./{app_path}",
            "-M", mode,
            "-m", model_path,
            "-p", prompt,
            "-H", str(height),
            "-W", str(width),
            "--steps", str(steps),
            "-s", str(seed),
            "--cfg-scale", str(cfg_scale),
            "--strength", str(strength),
            "--sampling-method", sampling_method,
            "--batch-count", str(batch_count),
            "--schedule", schedule,
            "--clip-skip", str(clip_skip),
        ]
        # Thêm tùy chọn Weight Type nếu được chỉ định
        if weight_type:
            command.extend(["--type", weight_type])
        
        # Thêm các tùy chọn khác
        if vae_tiling:
            command.append("--vae-tiling")
        if vae_on_cpu:
            command.append("--vae-on-cpu")
        if clip_on_cpu:
            command.append("--clip-on-cpu")
        if diffusion_fa:
            command.append("--diffusion-fa")
        if control_net_cpu:
            command.append("--control-net-cpu")
        if canny:
            command.append("--canny")
        if verbose:
            command.append("-v")
        # Thêm ảnh đầu vào nếu có
        if mode == "img2img" and init_image_path:
            command.extend(["-i", init_image_path])
        # Chạy lệnh và hiển thị log theo thời gian thực
        process = subprocess.Popen(
            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
        )
        logs = []
        for line in process.stdout:
            logs.append(line.strip())  # Lưu log vào danh sách
            print(line, end="")  # In log ra màn hình
        process.wait()  # Đợi tiến trình hoàn thành
        # Kiểm tra kết quả và trả về
        if process.returncode == 0:
            # Đường dẫn ảnh đầu ra mặc định
            output_path = "./output.png"
            if os.path.exists(output_path):
                # Tạo tên thư mục và tên file dựa trên ngày giờ hiện tại
                now = datetime.now()
                folder_name = f"SetImages+{now.strftime('%d+%m+%Y')}"
                file_name = f"Image+{now.strftime('%S+%M+%H')}+_{prompt[:50].replace(' ', '_')}.png"
                local_folder = os.path.join(".", folder_name)
                os.makedirs(local_folder, exist_ok=True)
                local_file_path = os.path.join(local_folder, file_name)
                # Di chuyển file đến thư mục mới
                os.rename(output_path, local_file_path)
                # Upload file lên Hugging Face Hub
                api = HfApi()
                repo_id = "ArrcttacsrjksX/Deffusion"
                remote_path = f"{folder_name}/{file_name}"
                upload_file(
                    path_or_fileobj=local_file_path,
                    path_in_repo=remote_path,
                    repo_id=repo_id,
                    token=hf_token,
                    repo_type="dataset",
                )
                print(f"Uploaded to Hugging Face: {remote_path}")
                return local_file_path, "\n".join(logs)
            else:
                return None, "\n".join(logs)
        else:
            error_log = process.stderr.read()  # Đọc lỗi
            logs.append(error_log)
            return None, "\n".join(logs)
    except Exception as e:
        return None, str(e)

# Giao diện Gradio
def toggle_image_input(mode):
    """Hiển thị hoặc ẩn ô Drop Image dựa trên mode."""
    return gr.update(visible=(mode == "img2img"))

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🌟 **Stable Diffusion Interface**
        Generate stunning images from text or modify existing images with AI-powered tools.
        """
    )
    # Thiết lập giao diện
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="🎨 Prompt", placeholder="Enter your creative idea here...", lines=2
            )
            mode = gr.Radio(
                choices=["txt2img", "img2img"], value="txt2img", label="Mode", interactive=True
            )
            init_image = gr.Image(
                label="Drop Image (for img2img mode)", type="pil", visible=False
            )
            mode.change(toggle_image_input, inputs=mode, outputs=init_image)
        with gr.Column():
            height = gr.Slider(
                128, 1024, value=512, step=64, label="Image Height (px)", interactive=True
            )
            width = gr.Slider(
                128, 1024, value=512, step=64, label="Image Width (px)", interactive=True
            )
            steps = gr.Slider(
                1, 100, value=20, step=1, label="Sampling Steps", interactive=True
            )
            seed = gr.Slider(
                -1, 10000, value=-1, step=1, label="Random Seed (-1 for random)", interactive=True
            )
            cfg_scale = gr.Slider(
                1, 20, value=7, step=0.1, label="CFG Scale", interactive=True
            )
            strength = gr.Slider(
                0, 1, value=0.75, step=0.01, label="Strength (img2img only)", interactive=True
            )
    with gr.Row():
        sampling_method = gr.Dropdown(
            choices=["euler", "euler_a", "heun", "dpm2", "dpm++2s_a", "dpm++2m", "dpm++2mv2", "ipndm", "ipndm_v", "lcm"],
            value="euler_a", label="Sampling Method", interactive=True
        )
        batch_count = gr.Slider(
            1, 10, value=1, step=1, label="Batch Count", interactive=True
        )
        schedule = gr.Dropdown(
            choices=["discrete", "karras", "exponential", "ays", "gits"],
            value="discrete", label="Denoiser Sigma Schedule", interactive=True
        )
    with gr.Row():
        clip_skip = gr.Slider(
            -1, 10, value=-1, step=1, label="CLIP Skip Layers", interactive=True
        )
        vae_tiling = gr.Checkbox(label="VAE Tiling", value=False)
        vae_on_cpu = gr.Checkbox(label="VAE on CPU", value=False)
        clip_on_cpu = gr.Checkbox(label="CLIP on CPU", value=False)
        diffusion_fa = gr.Checkbox(label="Diffusion Flash Attention", value=False)
        control_net_cpu = gr.Checkbox(label="ControlNet on CPU", value=False)
        canny = gr.Checkbox(label="Canny Preprocessor", value=False)
        verbose = gr.Checkbox(label="Verbose Logging", value=False)
    with gr.Row():
        weight_type = gr.Dropdown(
            choices=["f32", "f16", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2_K", "q3_K", "q4_K"],
            value=None,  # Mặc định không chọn
            label="Weight Type (Optional)",
            interactive=True
        )
    # Nút chạy và kết quả
    with gr.Row():
        run_button = gr.Button("🚀 Run", variant="primary")
    with gr.Row():
        output_image = gr.File(label="Download Image", interactive=False)
        log_output = gr.Textbox(label="Logs", interactive=False, lines=10)
    # Kết nối nút Run với hàm xử lý
    run_button.click(
        run_command_and_save,
        inputs=[
            prompt, mode, height, width, steps, seed, cfg_scale, strength, sampling_method,
            batch_count, schedule, clip_skip, vae_tiling, vae_on_cpu, clip_on_cpu, diffusion_fa,
            control_net_cpu, canny, verbose, init_image, weight_type  # Thêm weight_type
        ],
        outputs=[output_image, log_output],
    )

demo.launch()