app.py

import os
import math
import gradio as gr
import numpy as np
import torch
import safetensors.torch as sf

from PIL import Image
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler
from diffusers.models.attention_processor import AttnProcessor2_0
from transformers import CLIPTextModel, CLIPTokenizer
from briarmbg import BriaRMBG
from torch.hub import download_url_to_file
from spaces import ZeroGPU

# مقداردهی اولیه ZeroGPU
ZeroGPU.initialize()

# نصب کتابخانه‌های مورد نیاز در صورت اجرا روی Hugging Face Spaces
os.system("pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121")

# بررسی دسترسی به GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_name(0).startswith("A100") else torch.float16
if torch.cuda.is_available():
    print(f"✅ CUDA is available! GPU: {torch.cuda.get_device_name(0)}")
else:
    print("❌ No CUDA GPUs are available. Switching to CPU.")

# دانلود و بارگذاری مدل
model_path = './models/iclight_sd15_fc.safetensors'
if not os.path.exists(model_path):
    download_url_to_file(
        url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors',
        dst=model_path
    )

# بارگذاری مدل‌های مورد نیاز
sd15_name = 'stablediffusionapi/realistic-vision-v51'
tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder").to(device=device, dtype=dtype)
vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae").to(device=device, dtype=dtype)
unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet").to(device=device, dtype=dtype)
rmbg = BriaRMBG.from_pretrained("briaai/RMBG-1.4").to(device=device, dtype=torch.float32)

# تنظیم پردازشگر توجه برای بهینه‌سازی عملکرد
unet.set_attn_processor(AttnProcessor2_0())
vae.set_attn_processor(AttnProcessor2_0())

# تنظیم نمونه‌گیرها
scheduler = DPMSolverMultistepScheduler(
    num_train_timesteps=1000,
    beta_start=0.00085,
    beta_end=0.012,
    algorithm_type="sde-dpmsolver++",
    use_karras_sigmas=True,
    steps_offset=1
)

# ساخت لوله‌های پردازش تصویر
t2i_pipe = StableDiffusionPipeline.from_pretrained(
    sd15_name,
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=scheduler,
    safety_checker=None,
    requires_safety_checker=False,
    feature_extractor=None,
    image_encoder=None
).to(device)

i2i_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    sd15_name,
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=scheduler,
    safety_checker=None,
    requires_safety_checker=False,
    feature_extractor=None,
    image_encoder=None
).to(device)

# پاک کردن کش حافظه GPU برای جلوگیری از خطای Out Of Memory
torch.cuda.empty_cache()

# توابع پردازش تصویر
def process_image(prompt: str, width: int, height: int, num_steps: int, guidance: float):
    rng = torch.Generator(device=device).manual_seed(42)
    result = t2i_pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_steps,
        guidance_scale=guidance,
        generator=rng
    ).images[0]
    return result

# راه‌اندازی رابط کاربری با Gradio
block = gr.Blocks()
with block:
    with gr.Row():
        gr.Markdown("## IC-Light - Image Processing")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt")
            width = gr.Slider(256, 1024, value=512, step=64, label="Width")
            height = gr.Slider(256, 1024, value=640, step=64, label="Height")
            steps = gr.Slider(10, 100, value=50, step=5, label="Inference Steps")
            guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.5, label="Guidance Scale")
            submit = gr.Button("Generate Image")
        with gr.Column():
            result_image = gr.Image(type="pil", label="Generated Image")
    
    submit.click(fn=process_image, inputs=[prompt, width, height, steps, guidance], outputs=result_image)

block.launch(server_name='0.0.0.0')