import os from dotenv import load_dotenv load_dotenv() # Only needed locally # Prevent libgomp crashes in Spaces os.environ["OMP_NUM_THREADS"] = "1" os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" os.environ["HTTP_PROXY"] = "" os.environ["HTTPS_PROXY"] = "" os.environ["http_proxy"] = "" os.environ["https_proxy"] = "" import threading import requests from huggingface_hub import snapshot_download import gradio as gr import base64 import io import json import uuid # Import pose transfer related modules try: from diffusers import DDPMScheduler from defaults import pose_transfer_C as cfg from pose_transfer_train import build_model from models import UNet, VariationalAutoencoder import torch import numpy as np import pandas as pd from pose_utils import (cords_to_map, draw_pose_from_cords, load_pose_cords_from_strings) import random from PIL import Image from torchvision import transforms import copy except ImportError as e: print(f"Import error (models not downloaded yet): {e}") # ============================== # CONFIGURATION # ============================== MODEL_REPO = "recky101/new_l_cfld_model" PERSISTENT_DIR = "/data/models" # ============================== # GLOBALS # ============================== download_thread = None download_log = [] cancel_download = False model_ready = False vae = None model = None unet = None noise_scheduler = None test_pairs = None annotation_file = None # ============================== # UTILS # ============================== def log(msg: str): """Append a log message and return the full log as a string.""" global download_log download_log.append(msg) print(msg) return "\n".join(download_log) def test_huggingface_connectivity(): """Test if we can reach Hugging Face servers.""" try: response = requests.get("https://huggingface.co/", timeout=10) if response.status_code == 200: return True, "✅ Can reach Hugging Face" else: return False, f"❌ Hugging Face returned status {response.status_code}" except Exception as e: return False, f"❌ Cannot reach Hugging Face: {str(e)}" def is_model_ready(): """Check if required folders already exist.""" checkpoints = os.path.join(PERSISTENT_DIR, "checkpoints") pretrained = os.path.join(PERSISTENT_DIR, "pretrained_models") fashion = os.path.join(PERSISTENT_DIR, "fashion") return all(os.path.exists(path) for path in [checkpoints, pretrained, fashion]) def get_directory_tree(root_dir, indent=""): """Recursively build a directory tree string for logs.""" tree_str = "" try: items = sorted(os.listdir(root_dir)) except Exception as e: return f"{indent}❌ [Error accessing {root_dir}]: {e}\n" for i, item in enumerate(items): path = os.path.join(root_dir, item) connector = "└── " if i == len(items) - 1 else "├── " tree_str += f"{indent}{connector}{item}\n" if os.path.isdir(path): tree_str += get_directory_tree(path, indent + (" " if i == len(items) - 1 else "│ ")) return tree_str def download_individual_folders_with_retry(): """Alternative method with robust retry logic.""" try: folders_to_download = ["checkpoints", "pretrained_models", "fashion"] max_retries = 3 for folder in folders_to_download: if cancel_download: log("⛔ Download cancelled during individual folder download.") return folder_path = os.path.join(PERSISTENT_DIR, folder) if os.path.exists(folder_path): log(f"✅ Folder {folder} already exists, skipping...") continue for attempt in range(max_retries): try: log(f"⬇️ Downloading folder: {folder} (Attempt {attempt + 1}/{max_retries})") snapshot_download( repo_id=MODEL_REPO, local_dir=folder_path, resume_download=True, local_dir_use_symlinks=False, allow_patterns=f"{folder}/*" ) log(f"✅ Successfully downloaded {folder}") break except Exception as e: if attempt == max_retries - 1: log(f"❌ Failed to download {folder} after {max_retries} attempts: {str(e)}") else: log(f"⚠️ Attempt {attempt + 1} failed for {folder}: {str(e)}") import time time.sleep(5) log("✅ All folders processed.") except Exception as e: log(f"❌ Individual folder download failed: {str(e)}") def download_models(): """Download models with logging and show directory tree after completion.""" global cancel_download, model_ready, vae, model, unet, noise_scheduler, test_pairs, annotation_file try: # Test connectivity first success, message = test_huggingface_connectivity() log(message) if not success: log("🌐 Please check your internet connection and try again") return if is_model_ready(): log("✅ Models already downloaded. Skipping...") log("📂 Directory structure:") tree = get_directory_tree(PERSISTENT_DIR) log(f"\n{tree}") model_ready = True initialize_models() return os.makedirs(PERSISTENT_DIR, exist_ok=True) log("⬇️ Starting model download from Hugging Face Hub...") # Add retry logic max_retries = 3 for attempt in range(max_retries): try: if cancel_download: log("⛔ Download cancelled during process.") return log(f"🔄 Attempt {attempt + 1}/{max_retries}") snapshot_download( repo_id=MODEL_REPO, local_dir=PERSISTENT_DIR, resume_download=True, local_dir_use_symlinks=False, ) break except Exception as e: if attempt == max_retries - 1: raise e log(f"⚠️ Attempt {attempt + 1} failed: {str(e)}") import time time.sleep(10) if cancel_download: log("⛔ Download cancelled during process.") return log("✅ Download completed successfully.") log("📂 Listing downloaded directory structure...") tree = get_directory_tree(PERSISTENT_DIR) log(f"\n{tree}") model_ready = True initialize_models() except Exception as e: error_msg = str(e) log(f"❌ Download failed after {max_retries} attempts: {error_msg}") log("💡 Trying alternative download method...") download_individual_folders_with_retry() def initialize_models(): """Initialize the pose transfer models after download.""" global vae, model, unet, noise_scheduler, test_pairs, annotation_file try: log("🔄 Initializing pose transfer models...") # DEBUG: Print entire directory tree with sizes log("📂 Verifying model files in persistent directory...") for root, dirs, files in os.walk(PERSISTENT_DIR): level = root.replace(PERSISTENT_DIR, "").count(os.sep) indent = " " * 4 * (level) log(f"{indent}{os.path.basename(root)}/") subindent = " " * 4 * (level + 1) for f in files: size = os.path.getsize(os.path.join(root, f)) / (1024*1024) log(f"{subindent}{f} ({size:.2f} MB)") log("🔄 Initializing pose transfer models after tyeh checking ll thigns out...") try: # Initialize models noise_scheduler = DDPMScheduler.from_pretrained(os.path.join(PERSISTENT_DIR, "pretrained_models/scheduler/scheduler_config.json")) log("🔄 noise done") vae = VariationalAutoencoder(pretrained_path=os.path.join(PERSISTENT_DIR, "pretrained_models/vae")).eval().requires_grad_(False).cuda() log("🔄 vae done") model = build_model(cfg).eval().requires_grad_(False).cuda() log("🔄 model done") unet = UNet(cfg).eval().requires_grad_(False).cuda() log("🔄 unet done") except Exception as e: log(f"❌ Error during model initialization: {str(e)}") log(e) # Load model weights model.load_state_dict(torch.load( os.path.join(PERSISTENT_DIR, "checkpoints", "pytorch_model.bin"), map_location="cpu" ), strict=False) log("🔄 checkoitns done") unet.load_state_dict(torch.load( os.path.join(PERSISTENT_DIR, "checkpoints", "pytorch_model_1-001.bin"), map_location="cpu" ), strict=False) log("🔄 checkoitns22 done ") # Load test data test_pairs_path = os.path.join(PERSISTENT_DIR, "fashion", "fasion-resize-pairs-test.csv") test_pairs = pd.read_csv(test_pairs_path) annotation_file_path = os.path.join(PERSISTENT_DIR, "fashion", "fasion-resize-annotation-test.csv") annotation_file = pd.read_csv(annotation_file_path, sep=':') annotation_file = annotation_file.set_index('name') log("✅ Models initialized successfully") except Exception as e: log(f"❌ Error initializing models: {str(e)}") def build_pose_img(annotation_file, img_path): """Build pose image from annotation file.""" log(f"📄 img_path: {img_path}") log(f"📄 basename(img_path): {os.path.basename(img_path)}") log(f"📄 Index Sample: {annotation_file.index[:5]}") log(f"📄 Does key exist?: {os.path.basename(img_path) in annotation_file.index}") string = annotation_file.loc[os.path.basename(img_path)] array = load_pose_cords_from_strings(string['keypoints_y'], string['keypoints_x']) pose_map = torch.tensor(cords_to_map(array, (256, 256), (256, 176)).transpose(2, 0, 1), dtype=torch.float32) pose_img = torch.tensor(draw_pose_from_cords(array, (256, 256), (256, 176)).transpose(2, 0, 1) / 255., dtype=torch.float32) pose_img = torch.cat([pose_img, pose_map], dim=0) return pose_img def pose_transfer(source_image, test_pair_index): test_pair_index = int(test_pair_index) """Perform pose transfer from source image to target pose.""" global vae, model, unet, noise_scheduler, test_pairs, annotation_file if not model_ready: raise ValueError("Models not ready. Please download models first.") if test_pair_index < 0 or test_pair_index >= len(test_pairs): raise ValueError(f"Test pair index must be between 0 and {len(test_pairs)-1}") # Get target image path img_to_path = test_pairs.iloc[test_pair_index]["to"] log(f"🔄 img_to_path: {img_to_path}") # Build pose image pose_img_tensor = build_pose_img(annotation_file, img_to_path).unsqueeze(0) # Transform source image trans = transforms.Compose([ transforms.Resize([256, 256], interpolation=transforms.InterpolationMode.BICUBIC, antialias=True), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) img_from_tensor = trans(source_image).unsqueeze(0) # Perform pose transfer with torch.no_grad(): c_new, down_block_additional_residuals, up_block_additional_residuals = model({ "img_cond": img_from_tensor.cuda(), "pose_img": pose_img_tensor.cuda()}) noisy_latents = torch.randn((1, 4, 64, 64)).cuda() weight_dtype = torch.float32 bsz = 1 c_new = torch.cat([c_new[:bsz], c_new[:bsz], c_new[bsz:]]) down_block_additional_residuals = [torch.cat([torch.zeros_like(sample), sample, sample]).to(dtype=weight_dtype) for sample in down_block_additional_residuals] up_block_additional_residuals = {k: torch.cat([torch.zeros_like(v), torch.zeros_like(v), v]).to(dtype=weight_dtype) for k, v in up_block_additional_residuals.items()} noise_scheduler.set_timesteps(cfg.TEST.NUM_INFERENCE_STEPS) for t in noise_scheduler.timesteps: inputs = torch.cat([noisy_latents, noisy_latents, noisy_latents], dim=0) inputs = noise_scheduler.scale_model_input(inputs, timestep=t) noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new, down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals), up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals)) noise_pred_uc, noise_pred_down, noise_pred_full = noise_pred.chunk(3) noise_pred = noise_pred_uc + \ cfg.TEST.DOWN_BLOCK_GUIDANCE_SCALE * (noise_pred_down - noise_pred_uc) + \ cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_down) noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0] sampling_imgs = vae.decode(noisy_latents) * 0.5 + 0.5 # denormalize sampling_imgs = sampling_imgs.clamp(0, 1) # Convert to PIL image # output_img = Image.fromarray((sampling_imgs[0] * 255.).permute((1, 2, 0)).long().cpu().numpy().astype(np.uint8)).resize((256, 256)) # Convert tensor to PIL image without resizing output_img = Image.fromarray( (sampling_imgs[0] * 255.) .permute((1, 2, 0)) .long() .cpu() .numpy() .astype(np.uint8) ) # ✅ Save image in-memory as PNG, preserving original size img_bytes = io.BytesIO() output_img.save(img_bytes, format="PNG") img_bytes.seek(0) log("✅ Pose transfer completed successfully") log(f"🔄 output_img size: {output_img.size}") log(f"🔄 output_img butes: {img_bytes}") return output_img def start_download(): """Start model download in a separate thread.""" global download_thread, download_log, cancel_download if download_thread and download_thread.is_alive(): return "⚠️ Download already running..." download_log = [] cancel_download = False download_thread = threading.Thread(target=download_models) download_thread.start() return "📥 Download started..." def get_download_status(): """Get the latest log status.""" global download_thread status = "\n".join(download_log) if download_log else "Preparing download..." if download_thread and download_thread.is_alive(): return status + "\n\n⏳ Download in progress..." elif is_model_ready(): # tree = get_directory_tree(PERSISTENT_DIR) # log(f"\n{tree}") return status + "\n\n✅ Models are ready." return status def cancel_download_fn(): """Cancel request handler.""" global cancel_download cancel_download = True log("⛔ Download cancelled by user.") return "Download cancelled." # ============================== # GRADIO UI ONLY (NO API ENDPOINTS) # ============================== def gradio_pose_transfer(source_image, test_pair_index): test_pair_index = int(test_pair_index) """Gradio interface for pose transfer.""" try: if not model_ready: return None, "Models not ready. Please download models first." if test_pair_index < 0 or test_pair_index >= len(test_pairs): return None, f"Test pair index must be between 0 and {len(test_pairs)-1}" # Perform pose transfer output_image = pose_transfer(source_image, test_pair_index) # Save image to in-memory bytes (PNG) img_bytes = io.BytesIO() output_image.save(img_bytes, format="PNG") img_bytes.seek(0) return output_image, img_bytes, "Pose transfer successful" except Exception as e: import traceback tb = traceback.format_exc() log(f"❌ Error during pose transfer:\n{tb}") log(f"❌ the value recieved is as follow: \n{test_pair_index}") return None, f"Error during pose transfer: {str(e)}" with gr.Blocks() as demo: gr.Markdown("## 🧩 Model Downloader & Pose Transfer") gr.Markdown(f"**Model Source:** [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO})") gr.Markdown("**Required folders:** checkpoints/, pretrained_models/, fashion/") with gr.Tab("Model Download"): with gr.Row(): start_btn = gr.Button("📥 Download Models") cancel_btn = gr.Button("❌ Cancel Download") status_box = gr.Textbox( label="Download Logs", lines=25, interactive=False, placeholder="Click 'Download Models' to start downloading..." ) # Button bindings start_btn.click(fn=start_download, inputs=None, outputs=status_box) cancel_btn.click(fn=cancel_download_fn, inputs=None, outputs=status_box) # Periodic refresh of logs demo.load(fn=get_download_status, inputs=None, outputs=status_box, every=2) with gr.Tab("Pose Transfer"): gr.Markdown("## Pose Transfer") with gr.Row(): with gr.Column(): source_image = gr.Image(label="Source Image", type="pil") test_pair_index = gr.Number( label="Test Pair Index", value=0, minimum=0, maximum=4039 ) generate_btn = gr.Button("🚀 Generate Pose Transfer") with gr.Column(): output_image = gr.Image(label="Output Image", type="pil") download_btn = gr.File(label="Download Output", file_types=[".png"]) status_message = gr.Textbox(label="Status", interactive=False) generate_btn.click( fn=gradio_pose_transfer, inputs=[source_image, test_pair_index], outputs=[output_image, status_message] ) # ============================== # START APP # ============================== from fastapi import FastAPI, Request from fastapi.responses import JSONResponse import traceback def read_image_from_input(image_input: str): """ Accepts either: - Base64 image string ("data:image/png;base64,...") - Image URL ("https://...") Returns: PIL.Image """ try: if image_input.startswith("http://") or image_input.startswith("https://"): response = requests.get(image_input) response.raise_for_status() return Image.open(io.BytesIO(response.content)).convert("RGB") else: # Base64 input if "," in image_input: image_input = image_input.split(",", 1)[1] image_bytes = base64.b64decode(image_input) return Image.open(io.BytesIO(image_bytes)).convert("RGB") except Exception as e: raise ValueError(f"Invalid image input: {str(e)}") # from fastapi import FastAPI, Form, UploadFile # from fastapi.responses import JSONResponse # from fastapi.middleware.cors import CORSMiddleware # from PIL import Image # import base64 # from io import BytesIO # import uvicorn # app = FastAPI() # # Allow CORS (needed for Postman / frontends) # app.add_middleware( # CORSMiddleware, # allow_origins=["*"], # allow_credentials=True, # allow_methods=["*"], # allow_headers=["*"], # ) # @app.post("/pose-transfer") # async def pose_transfer_api( # test_pair_index: int = Form(...), # source_image: UploadFile = None, # image_base64: str = Form(None) # ): # try: # # Get input image # if source_image: # image_bytes = await source_image.read() # image = Image.open(BytesIO(image_bytes)).convert("RGB") # elif image_base64: # if "," in image_base64: # image_base64 = image_base64.split(",", 1)[1] # image_bytes = base64.b64decode(image_base64) # image = Image.open(BytesIO(image_bytes)).convert("RGB") # else: # return JSONResponse( # content={"status": "error", "message": "No image provided"}, # status_code=400 # ) # # 🔹 Call your pose transfer model # output_image = pose_transfer(image, test_pair_index) # # Convert result to base64 # buffer = BytesIO() # output_image.save(buffer, format="PNG") # img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") # return JSONResponse( # content={ # "status": "success", # "message": "Pose transfer completed", # "output_image_base64": f"data:image/png;base64,{img_base64}" # } # ) # except Exception as e: # import traceback # tb = traceback.format_exc() # return JSONResponse( # content={"status": "error", "message": str(e), "traceback": tb}, # status_code=500 # ) # if __name__ == "__main__": # if is_model_ready(): # model_ready = True # initialize_models() # uvicorn.run(app, host="0.0.0.0", port=7860) from fastapi import FastAPI, Form, UploadFile from fastapi.responses import JSONResponse from fastapi.middleware.cors import CORSMiddleware from PIL import Image import base64 from io import BytesIO import uvicorn import cloudinary import cloudinary.uploader import os # Configure Cloudinary from environment variables cloudinary.config( cloud_name=os.environ.get("CLOUDINARY_CLOUD_NAME"), api_key=os.environ.get("CLOUDINARY_API_KEY"), api_secret=os.environ.get("CLOUDINARY_API_SECRET") ) app = FastAPI() # Allow CORS (needed for Postman / frontends) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @app.post("/pose-transfer") async def pose_transfer_api( test_pair_index: int = Form(...), source_image: UploadFile = None, image_base64: str = Form(None) ): try: # Get input image if source_image: image_bytes = await source_image.read() image = Image.open(BytesIO(image_bytes)).convert("RGB") elif image_base64: if "," in image_base64: image_base64 = image_base64.split(",", 1)[1] image_bytes = base64.b64decode(image_base64) image = Image.open(BytesIO(image_bytes)).convert("RGB") else: return JSONResponse( content={"status": "error", "message": "No image provided"}, status_code=400 ) # Check if Cloudinary is configured if not all([cloudinary.config().cloud_name, cloudinary.config().api_key, cloudinary.config().api_secret]): return JSONResponse( content={"status": "error", "message": "Cloudinary not properly configured"}, status_code=500 ) # 🔹 Call your pose transfer model output_image = pose_transfer(image, test_pair_index) # Save image to buffer buffer = BytesIO() output_image.save(buffer, format="PNG") buffer.seek(0) # Upload to Cloudinary upload_result = cloudinary.uploader.upload( buffer, folder="pose_transfer", # Optional folder in Cloudinary public_id=f"pose_transfer_{uuid.uuid4().hex[:8]}", # Unique ID overwrite=True, resource_type="image" ) # Get the URL from Cloudinary response image_url = upload_result.get('secure_url', upload_result.get('url')) return JSONResponse( content={ "status": "success", "message": "Pose transfer completed and image uploaded to Cloudinary", "image_url": image_url } ) except Exception as e: import traceback tb = traceback.format_exc() return JSONResponse( content={"status": "error", "message": str(e), "traceback": tb}, status_code=500 ) if __name__ == "__main__": if is_model_ready(): model_ready = True initialize_models() uvicorn.run(app, host="0.0.0.0", port=7860) # if __name__ == "__main__": # # Initialize models if they exist # if is_model_ready(): # model_ready = True # initialize_models() # demo.launch(server_name="0.0.0.0", server_port=7860)