Spaces:
Running
on
Zero
Running
on
Zero
| 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) --- | |
| 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() |