prev_app.py

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)