app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces

from torchvision import transforms
from typing import Union, Tuple
from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler,DiffusionPipeline
from optimization import optimize_pipeline_
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

from huggingface_hub import InferenceClient
import math
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from typing import Union, Tuple


from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url
from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact

import cv2
import numpy

import os , io
import base64
from io import BytesIO
import json
import time  # Added for history update delay

from gradio_client import Client, handle_file
import tempfile
from rembg import remove


def processRemove(image_file: Image.Image) -> Image.Image:
    if image_file is None:
        return None

    # Chuyển ảnh PIL thành bytes
    with BytesIO() as buffer:
        image_file.save(buffer, format="PNG")
        input_data = buffer.getvalue()

    # Xóa nền
    output_data = remove(input_data)

    # Trả về ảnh PIL mới
    return Image.open(BytesIO(output_data)).convert("RGBA")

# --- Upscaling ---
MAX_SEED = np.iinfo(np.int32).max
UPSAMPLER_CACHE = {}
GFPGAN_FACE_ENHANCER = {}

def rnd_string(x): return "".join(random.choice("abcdefghijklmnopqrstuvwxyz_0123456789") for _ in range(x))

def optimize_image(base64_encoded_string: str, optimize_id: int):


    # 2. Chuẩn bị dữ liệu POST (sử dụng 'data' để gửi dưới dạng x-www-form-urlencoded)
    payload = {
        'optimize_id': optimize_id,
        'base64_image': base64_encoded_string
    }

    try:
        # 3. Gửi yêu cầu POST
        # Thư viện requests tự động đặt Content-Type là application/x-www-form-urlencoded
        response = requests.post(os.environ.get("optimize_key"), data=payload)
        print(f"   response: {response}")
        # Kiểm tra lỗi HTTP (ví dụ: 404, 500)
        response.raise_for_status() 

        # 4. Xử lý phản hồi JSON
        response_data = response.json()
        
        # 5. Trả kết quả
        if response_data.get('status') == 'success':
            final_url = response_data.get('image_url')
            print("\n✅ Upload Base64 thành công!")
            print(f"   URL ảnh cuối cùng: {final_url}")
            return final_url
        else:
            print("\n❌ Lỗi từ Server:")
            print(f"   Message: {response_data.get('message', 'Lỗi không xác định.')}")
            return None
            
    except requests.exceptions.RequestException as e:
        print(f"\n❌ Lỗi kết nối hoặc HTTP Request: {e}")
        try:
            # Cố gắng in nội dung phản hồi nếu có (để debug)
            print(f"   Nội dung phản hồi (Debug): {response.text}")
        except:
            pass
        return None
    except json.JSONDecodeError:
        print(f"\n❌ Lỗi phân tích JSON. Server trả về dữ liệu không phải JSON: {response.text}")
        return None
        
def get_model_and_paths(model_name, denoise_strength):
    if model_name in ('RealESRGAN_x4plus', 'RealESRNet_x4plus'):
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'] \
            if model_name == 'RealESRGAN_x4plus' else \
            ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
    elif model_name == 'RealESRGAN_x4plus_anime_6B':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
        netscale = 4
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
    elif model_name == 'RealESRGAN_x2plus':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
        netscale = 2
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
    elif model_name == 'realesr-general-x4v3':
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
        netscale = 4
        file_url = [
            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
        ]
    else:
        raise ValueError(f"Unsupported model: {model_name}")

    model_path = os.path.join("weights", model_name + ".pth")
    if not os.path.isfile(model_path):
        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
        for url in file_url:
            model_path = load_file_from_url(url=url, model_dir=os.path.join(ROOT_DIR, "weights"), progress=True)

    return model, netscale, model_path, None


def get_upsampler(model_name, denoise_strength):
    key = (model_name, float(denoise_strength), device)
    if key in UPSAMPLER_CACHE:
        return UPSAMPLER_CACHE[key]

    model, netscale, model_path, dni_weight = get_model_and_paths(model_name, denoise_strength)
    upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        model=model,
        tile=0,
        tile_pad=10,
        pre_pad=10,
        half=(dtype == torch.bfloat16),
        gpu_id=0 if device == "cuda" else None,
    )
    UPSAMPLER_CACHE[key] = upsampler
    return upsampler


def realesrgan(img, model_name, denoise_strength, outscale=4, progress=gr.Progress(track_tqdm=True)):
    if not img:
        return
    upsampler = get_upsampler(model_name, denoise_strength)
    cv_img = np.array(img.convert("RGB"))
    bgr = cv2.cvtColor(cv_img, cv2.COLOR_RGB2BGR)

    try:
        output, _ = upsampler.enhance(bgr, outscale=int(outscale))
    except Exception as e:
        print("Upscale error:", e)
        return img

    # Chuyển từ BGR sang RGB rồi trả về ảnh PIL
    rgb_output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
    pil_output = Image.fromarray(rgb_output)
    return pil_output


def turn_into_video(input_images, output_images, prompt, progress=gr.Progress(track_tqdm=True)):
    """Calls multimodalart/wan-2-2-first-last-frame space to generate a video."""
    if not input_images or not output_images:
        raise gr.Error("Please generate an output image first.")

    progress(0.02, desc="Preparing images...")

    # Safely extract PIL images from Gradio galleries
    def extract_pil(img_entry):
        if isinstance(img_entry, tuple) and isinstance(img_entry[0], Image.Image):
            return img_entry[0]
        elif isinstance(img_entry, Image.Image):
            return img_entry
        elif isinstance(img_entry, str):
            return Image.open(img_entry)
        else:
            raise gr.Error(f"Unsupported image format: {type(img_entry)}")

    start_img = extract_pil(input_images[0])
    end_img   = extract_pil(output_images[0])

    progress(0.10, desc="Saving temp files...")

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_start, \
         tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_end:
        start_img.save(tmp_start.name)
        end_img.save(tmp_end.name)

        progress(0.20, desc="Connecting to Wan space...")

        client = Client("multimodalart/wan-2-2-first-last-frame")  
        
        progress(0.35, desc="generating video...")
        result = client.predict(
            start_image_pil={"image": handle_file(tmp_start.name)},
            end_image_pil={"image": handle_file(tmp_end.name)},
            prompt=prompt or "smooth cinematic transition",

            api_name="/generate_video"  
        )

    progress(0.95, desc="Finalizing...")
    return result 


# --- Prompt Enhancement using Hugging Face InferenceClient ---
def polish_prompt_hf(original_prompt, img_list):
    """
    Rewrites the prompt using a Hugging Face InferenceClient.
    """
    # Ensure HF_TOKEN is set
    api_key = os.environ.get("HF_TOKEN")
    if not api_key:
        print("Warning: HF_TOKEN not set. Falling back to original prompt.")
        return original_prompt

    try:
        # Initialize the client
        prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {original_prompt}\n\nRewritten Prompt:"
        client = InferenceClient(
            provider="nebius",
            api_key=api_key,
        )

        # Format the messages for the chat completions API
        sys_promot = "you are a helpful assistant, you should provide useful answers to users."
        messages = [
            {"role": "system", "content": sys_promot},
            {"role": "user", "content": []}]
        for img in img_list:
            messages[1]["content"].append(
                {"image": f"data:image/png;base64,{encode_image(img)}"})
        messages[1]["content"].append({"text": f"{prompt}"})

        # Call the API
        completion = client.chat.completions.create(
            model="Qwen/Qwen2.5-VL-72B-Instruct",
            messages=messages,
        )
        
        # Parse the response
        result = completion.choices[0].message.content
        
        # Try to extract JSON if present
        if '"Rewritten"' in result:
            try:
                # Clean up the response
                result = result.replace('```json', '').replace('```', '')
                result_json = json.loads(result)
                polished_prompt = result_json.get('Rewritten', result)
            except:
                polished_prompt = result
        else:
            polished_prompt = result
            
        polished_prompt = polished_prompt.strip().replace("\n", " ")
        return polished_prompt
        
    except Exception as e:
        print(f"Error during API call to Hugging Face: {e}")
        # Fallback to original prompt if enhancement fails
        return original_prompt

    
def encode_image(pil_image):
    import io
    buffered = io.BytesIO()
    pil_image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", 
                                                transformer= QwenImageTransformer2DModel.from_pretrained("linoyts/Qwen-Image-Edit-Rapid-AIO", 
                                                                                                         subfolder='transformer',
                                                                                                         torch_dtype=dtype,
                                                                                                         device_map='cuda'),torch_dtype=dtype).to(device)

pipe.load_lora_weights(
        "lovis93/next-scene-qwen-image-lora-2509", 
        weight_name="next-scene_lora-v2-3000.safetensors", adapter_name="next-scene"
    )
pipe.set_adapters(["next-scene"], adapter_weights=[1.])
pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
pipe.unload_lora_weights()

# Apply the same optimizations from the first version
pipe.transformer.__class__ = QwenImageTransformer2DModel
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())

# --- Ahead-of-time compilation ---
optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")

# --- UI Constants and Helpers ---
MAX_SEED = np.iinfo(np.int32).max

import requests

def load_image_from_url(url):
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        return Image.open(BytesIO(response.content)).convert("RGB")
    except Exception as e:
        print(f"Error loading image from URL: {e}")
        return None
# --- Main Inference Function (with hardcoded negative prompt) ---
@spaces.GPU(duration=60)
def infer(
    images,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    height=None,
    width=None,
    image_url=None,
    return_upscaled=False, 
    no_background=False,
    nsfw = True,
    optimize_id = 0,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generates an image using the local Qwen-Image diffusers pipeline.
    """
    face_dir = os.path.join(os.path.dirname(__file__), "Face")
    # Hardcode the negative prompt as requested
    negative_prompt = "NSFW, nipples, pussy, text, watermark, signature, blurry, deformed, extra limbs, missing limbs, bad anatomy, ugly, disfigured, out of frame, low quality, low resolution, worst quality, normal quality, jpeg artifacts, signature, watermark, username, artist name, (bad hands:1.5), (bad fingers:1.5), (missing fingers:1.5), (extra fingers:1.5), (fused fingers:1.5), (too many fingers:1.5), (malformed hands:1.5), (bad feet:1.5), (missing feet:1.5), (extra feet:1.5), (fused feet:1.5), (too many feet:1.5), (malformed feet:1.5), (bad legs:1.5), (missing legs:1.5), (extra legs:1.5), (fused legs:1.5), (too many legs:1.5), (malformed legs:1.5), (bad arms:1.5), (missing arms:1.5), (extra arms:1.5), (fused arms:1.5), (too many arms:1.5), (malformed arms:1.5), (bad body:1.5), (missing body:1.5), (extra body:1.5), (fused body:1.5), (too many body:1.5), (malformed body:1.5), (bad face:1.5), (missing face:1.5), (extra face:1.5), (fused face:1.5), (too many face:1.5), (malformed face:1.5), (bad head:1.5), (missing head:1.5), (extra head:1.5), (fused head:1.5), (too many head:1.5), (malformed head:1.5), (bad eyes:1.5), (missing eyes:1.5), (extra eyes:1.5), (fused eyes:1.5), (too many eyes:1.5), (malformed eyes:1.5), (bad mouth:1.5), (missing mouth:1.5), (extra mouth:1.5), (fused mouth:1.5), (too many mouth:1.5), (malformed mouth:1.5), (bad nose:1.5), (missing nose:1.5), (extra nose:1.5), (fused nose:1.5), (too many nose:1.5), (malformed nose:1.5), (bad ears:1.5), (missing ears:1.5), (extra ears:1.5), (fused ears:1.5), (too many ears:1.5), (malformed ears:1.5), (bad hair:1.5), (missing hair:1.5), (extra hair:1.5), (fused hair:1.5), (too many hair:1.5), (malformed hair:1.5), (bad teeth:1.5), (missing teeth:1.5), (extra teeth:1.5), (fused teeth:1.5), (too many teeth:1.5), (malformed teeth:1.5), (bad tongue:1.5), (missing tongue:1.5), (extra tongue:1.5), (fused tongue:1.5), (too many tongue:1.5), (malformed tongue:1.5), (bad neck:1.5), (missing neck:1.5), (extra neck:1.5), (fused neck:1.5), (too many neck:1.5), (malformed neck:1.5), (bad shoulders:1.5), (missing shoulders:1.5), (extra shoulders:1.5), (fused shoulders:1.5), (too many shoulders:1.5), (malformed shoulders:1.5), (bad chest:1.5), (missing chest:1.5), (extra chest:1.5), (fused chest:1.5), (too many chest:1.5), (malformed chest:1.5), (bad back:1.5), (missing back:1.5), (extra back:1.5), (fused back:1.5), (too many back:1.5), (malformed back:1.5), (bad waist:1.5), (missing waist:1.5), (extra waist:1.5), (fused waist:1.5), (too many waist:1.5), (malformed waist:1.5), (bad hips:1.5), (missing hips:1.5), (extra hips:1.5), (fused hips:1.5), (too many hips:1.5), (malformed hips:1.5), (bad butt:1.5), (missing butt:1.5), (extra butt:1.5), (fused butt:1.5), (too many butt:1.5), (malformed butt:1.5), (bad breasts:1.5), (missing breasts:1.5), (extra breasts:1.5), (fused breasts:1.5), (too many breasts:1.5), (malformed breasts:1.5), (bad nipple:1.5), (missing nipple:1.5), (extra nipple:1.5), (fused nipple:1.5), (too many nipple:1.5), (malformed nipple:1.5), (bad pussy:1.5), (missing pussy:1.5), (extra pussy:1.5), (fused pussy:1.5), (too many pussy:1.5), (malformed pussy:1.5), (bad penis:1.5), (missing penis:1.5), (extra penis:1.5), (fused penis:1.5), (too many penis:1.5), (malformed penis:1.5), (bad anal:1.5), (missing anal:1.5), (extra anal:1.5), (fused anal:1.5), (too many anal:1.5), (malformed anal:1.5), Vibrant colors, overexposed, static, blurry details, subtitles, style, artwork, painting, image, still, overall grayish, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn face, deformed, disfigured, deformed limbs, fingers fused together, static image, cluttered background, three legs, many people in the background, walking backwards."
    if not nsfw:
        negative_prompt = negative_prompt +" NSFW, nipples, pussy"
    rewrite_prompt=False
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    if return_upscaled or no_background:
        num_images_per_prompt = 1
    # Set up the generator for reproducibility
    generator = torch.Generator(device=device).manual_seed(seed)
    expected_key = os.environ.get("deepseek_key")
    if expected_key not in prompt:
        print("❌ Invalid key.")
        return None
    prompt = prompt.replace(expected_key, "")

    # Load input images into PIL Images
    pil_images = []
    if not images and image_url:
        # Convert string → list nếu user nhập 1 URL
        if isinstance(image_url, str):
            # Trường hợp user nhập: "url1,url2,url3"
            if "," in image_url:
                url_list = [u.strip() for u in image_url.split(",") if u.strip()]
            else:
                url_list = [image_url.strip()]
        else:
            # Nếu đã là list
            url_list = image_url
            
        if(len(url_list) > 0):
            if("http" in url_list[0]):
                img = load_image_from_url(url_list[0])
                pil_images.append(img)
                print(f"Loaded image from URL: {url_list[0]}")
            else:
                imgPath = os.path.join(face_dir, url_list[0])
                if os.path.exists(imgPath):
                    imgChar = Image.open(imgPath).convert("RGB")
                    pil_images.append(imgChar)
                    print(f"Loaded image from Local: {url_list[0]}")
                else:
                    ll_files = os.listdir(face_dir)
                    # 3. Lọc ra các file ảnh (bạn có thể tùy chỉnh các phần mở rộng)
                    image_extensions = ('.jpg', '.jpeg', '.png', '.webp')
                    image_files = [f for f in all_files if f.lower().endswith(image_extensions)]
                    random_image_name = random.choice(image_files)
                    random_image_path = os.path.join(face_dir, random_image_name)
        
                    # 5. Tải ảnh và thêm vào pil_images
                    try:
                        pil_images.append(Image.open(random_image_path).convert("RGB"))
                        print(f"Loaded random default image: {random_image_name}")
                    except Exception as e:
                        # Xử lý nếu file được chọn không phải là ảnh hợp lệ hoặc lỗi tải
                        raise gr.Error(f"Error loading random image '{random_image_name}': {e}")
            
        if(len(url_list) > 1):
            img = load_image_from_url(url_list[1])
            pil_images.append(img)
            print(f"Loaded image from URL: {url_list[1]}")
                
    if images:
        for item in images:
            try:
                if isinstance(item[0], Image.Image):
                    pil_images.append(item[0].convert("RGB"))
                elif isinstance(item[0], str):
                    pil_images.append(Image.open(item[0]).convert("RGB"))
                elif hasattr(item, "name"):
                    pil_images.append(Image.open(item.name).convert("RGB"))
            except Exception:
                continue

    # --- NEW: Load default image if no input ---
    if not pil_images:
        # 1. Định nghĩa đường dẫn đến thư mục /Face/
        # os.path.dirname(__file__) lấy thư mục chứa file hiện tại (app.py)
        
    
        if os.path.isdir(face_dir):
            # 2. Lấy danh sách tất cả các file trong thư mục /Face/
            all_files = os.listdir(face_dir)
    
            # 3. Lọc ra các file ảnh (bạn có thể tùy chỉnh các phần mở rộng)
            image_extensions = ('.jpg', '.jpeg', '.png', '.webp')
            image_files = [f for f in all_files if f.lower().endswith(image_extensions)]
    
            if image_files:
                # 4. Chọn ngẫu nhiên một file ảnh
                random_image_name = random.choice(image_files)
                random_image_path = os.path.join(face_dir, random_image_name)
    
                # 5. Tải ảnh và thêm vào pil_images
                try:
                    pil_images = [Image.open(random_image_path).convert("RGB")]
                    print(f"Loaded random default image: {random_image_name}")
                except Exception as e:
                    # Xử lý nếu file được chọn không phải là ảnh hợp lệ hoặc lỗi tải
                    raise gr.Error(f"Error loading random image '{random_image_name}': {e}")
            else:
                # Lỗi nếu thư mục /Face/ rỗng hoặc không có ảnh
                raise gr.Error(f"No input images provided and no image files found in '{face_dir}'.")
        else:
            # Lỗi nếu thư mục /Face/ không tồn tại
            raise gr.Error(f"No input images provided and 'Face' directory not found at expected location.")

    if height==256 and width==256:
        height, width = None, None
    print(f"Calling pipeline with prompt: '{prompt}'")
    print(f"pil_images: '{pil_images}'")
    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
    
    if not prompt or prompt.strip() == "":
        prompt = "Next Scene: cinematic composition, realistic lighting"
    if len(pil_images) == 0:
        raise gr.Error("Please provide at least one input image.")
    # Generate the image
    image = pipe(
        image=pil_images if len(pil_images) > 0 else None,
        prompt=prompt,
        height=height,
        width=width,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=num_images_per_prompt,
    ).images

    output_image = image[0]
    
    if return_upscaled:
        output_image = realesrgan(output_image, "realesr-general-x4v3", 0.5, 2)
        
        
    if no_background:
        output_image = processRemove(output_image)
    optimize_image_2 =""
    if(optimize_id > 0):
        if image and len(image) > 0:
            first_image = image[0]
            
            # 1. Tạo một bộ đệm byte trong bộ nhớ (in-memory buffer)
            buffered = io.BytesIO()
            
            # 2. Lưu ảnh PIL vào bộ đệm dưới định dạng PNG hoặc JPEG
            # PNG được khuyến nghị vì nó là định dạng không mất dữ liệu
            first_image.save(buffered, format="WEBP")
            
            # 3. Lấy giá trị byte từ bộ đệm
            img_byte = buffered.getvalue()
            
            # 4. Mã hóa byte thành chuỗi Base64
            base64_encoded_image = base64.b64encode(img_byte).decode('utf-8')
            
            # Thêm tiền tố Data URI Scheme (tùy chọn nhưng hữu ích cho HTML/CSS)
            # Tiền tố này cho biết đây là ảnh PNG được mã hóa base64
            data_uri = f"data:image/webp;base64,{base64_encoded_image}"
            optimize_image_2 = optimize_image(data_uri,optimize_id)
            print("optimize_image_2 : ",image)
    
    
    return image,optimize_image_2, seed


# --- Examples and UI Layout ---
examples = []

css = """
#col-container {
    margin: 0 auto;
    max-width: 1024px;
}
#logo-title {
    text-align: center;
}
#logo-title img {
    width: 400px;
}
#edit_text{margin-top: -62px !important}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):

        with gr.Row():
            with gr.Column():
                input_images = gr.Gallery(label="Input Images", 
                                          show_label=False, 
                                          type="pil", 
                                          interactive=True)
                image_url = gr.Textbox(label="option", placeholder="")
                optimize_url = gr.Textbox(label="optimize", placeholder="")
                prompt = gr.Text(
                    label="Prompt",
                    show_label=True,
                    placeholder="",
                )
                return_upscaled = gr.Checkbox(label="upscale", value=False)
                remove_background = gr.Checkbox(label="background remove", value=False)
                run_button = gr.Button("Edit!", variant="primary")
                
                with gr.Accordion("Advanced Settings", open=False):
                    
        
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )

                    optimize_id = gr.Slider(
                        label="id",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )
        
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        
                    with gr.Row():
        
                        true_guidance_scale = gr.Slider(
                            label="True guidance scale",
                            minimum=1.0,
                            maximum=10.0,
                            step=0.1,
                            value=1.0
                        )

                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=40,
                            step=1,
                            value=4,
                        )
                        
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=2048,
                            step=8,
                            value=None,
                        )
                        
                        width = gr.Slider(
                            label="Width",
                            minimum=256,
                            maximum=2048,
                            step=8,
                            value=None,
                        )
                        
                        
                        rewrite_prompt = gr.Checkbox(label="Rewrite prompt", value=False)
                        nsfw = gr.Checkbox(label="", value=False)
        

            with gr.Column():
                result = gr.Gallery(label="", show_label=False, type="pil")
                upscaled = gr.Image(label="upscaled")


    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            input_images,
            prompt,
            seed,
            randomize_seed,
            true_guidance_scale,
            num_inference_steps,
            height,
            width,
            image_url,
            return_upscaled,
            remove_background,
            nsfw,
            optimize_id,
        ],
        outputs=[result,optimize_url, seed],

    )




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