app.py

import requests
import os
import gradio as gr
from huggingface_hub import update_repo_visibility, whoami, upload_folder, create_repo, upload_file # Removed duplicate update_repo_visibility
from slugify import slugify
# import gradio as gr # Already imported
import re
import uuid
from typing import Optional, Dict, Any
import json
# from bs4 import BeautifulSoup # Not used

TRUSTED_UPLOADERS = ["AI_Art_Factory", "KappaNeuro", "CiroN2022", "Norod78", "joachimsallstrom", "blink7630", "e-n-v-y", "DoctorDiffusion", "RalFinger", "artificialguybr"]

# --- Model Mappings ---
MODEL_MAPPING_IMAGE = {
    "SDXL 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
    "SDXL 0.9": "stabilityai/stable-diffusion-xl-base-1.0", # Usually mapped to 1.0
    "SD 1.5": "runwayml/stable-diffusion-v1-5",
    "SD 1.4": "CompVis/stable-diffusion-v1-4",
    "SD 2.1": "stabilityai/stable-diffusion-2-1-base",
    "SD 2.0": "stabilityai/stable-diffusion-2-base",
    "SD 2.1 768": "stabilityai/stable-diffusion-2-1",
    "SD 2.0 768": "stabilityai/stable-diffusion-2",
    "SD 3": "stabilityai/stable-diffusion-3-medium-diffusers", # Assuming medium, adjust if others are common
    "SD 3.5": "stabilityai/stable-diffusion-3.5-large", # Assuming large, adjust
    "SD 3.5 Large": "stabilityai/stable-diffusion-3.5-large",
    "SD 3.5 Medium": "stabilityai/stable-diffusion-3.5-medium",
    "SD 3.5 Large Turbo": "stabilityai/stable-diffusion-3.5-large-turbo",
    "Flux.1 D": "black-forest-labs/FLUX.1-dev",
    "Flux.1 S": "black-forest-labs/FLUX.1-schnell",
    "Pony": "AiAF/Diffusers_Pony-Diffusion-V6",
}

MODEL_MAPPING_VIDEO = {
    "LTXV": "Lightricks/LTX-Video-0.9.7-dev",
    "Wan Video 1.3B t2v": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
    "Wan Video 14B t2v": "Wan-AI/Wan2.1-T2V-14B-Diffusers",
    "Wan Video 14B i2v 480p": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
    "Wan Video 14B i2v 720p": "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
    "Hunyuan Video": "hunyuanvideo-community/HunyuanVideo-I2V", # Default, will be overridden by choice
}

SUPPORTED_CIVITAI_BASE_MODELS = list(MODEL_MAPPING_IMAGE.keys()) + list(MODEL_MAPPING_VIDEO.keys())


cookie_info = os.environ.get("COOKIE_INFO")

headers = {
        "authority": "civitai.com",
        "accept": "*/*",
        "accept-language": "en-US,en;q=0.9", # Simplified
        "content-type": "application/json",
        "cookie": cookie_info, # Use the env var
        "sec-ch-ua": "\"Chromium\";v=\"118\", \"Not_A Brand\";v=\"99\"", # Example, update if needed
        "sec-ch-ua-mobile": "?0",
        "sec-ch-ua-platform": "\"Windows\"", # Example
        "sec-fetch-dest": "empty",
        "sec-fetch-mode": "cors",
        "sec-fetch-site": "same-origin",
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/118.0.0.0 Safari/537.36" # Example
}

def get_json_data(url):
    url_split = url.split('/')
    if len(url_split) < 5 or not url_split[4].isdigit():
        print(f"Invalid Civitai URL format or model ID not found: {url}")
        gr.Warning(f"Invalid Civitai URL format. Ensure it's like 'https://civitai.com/models/YOUR_MODEL_ID/MODEL_NAME'. Problem with: {url}")
        return None
    api_url = f"https://civitai.com/api/v1/models/{url_split[4]}"
    try:
        response = requests.get(api_url)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error fetching JSON data from {api_url}: {e}")
        gr.Warning(f"Error fetching data from Civitai API for {url_split[4]}: {e}")
        return None

def check_nsfw(json_data: Dict[str, Any]) -> bool:
    """
    Returns True if the model or any of its versions/images are NSFW.
    We no longer block NSFW, only tag it.
    """
    nsfw_flag = False

    if json_data.get("nsfw", False):
        nsfw_flag = True

    if json_data.get("nsfwLevel", 0) > 0:
        nsfw_flag = True

    for model_version in json_data.get("modelVersions", []):
        if model_version.get("nsfwLevel", 0) > 0:
            nsfw_flag = True
        for image_obj in model_version.get("images", []):
            if image_obj.get("nsfwLevel", 0) > 0:
                nsfw_flag = True

    return nsfw_flag



def get_prompts_from_image(image_id_str: str):
    # image_id_str could be non-numeric if URL parsing failed or format changed
    try:
        image_id = int(image_id_str)
    except ValueError:
        print(f"Invalid image_id_str for TRPC call: {image_id_str}. Skipping prompt fetch.")
        return "", ""

    print(f"Fetching prompts for image_id: {image_id}")
    url = f'https://civitai.com/api/trpc/image.getGenerationData?input={{"json":{{"id":{image_id}}}}}'
    
    prompt = ""
    negative_prompt = ""
    try:
        response = requests.get(url, headers=headers, timeout=10) # Added timeout
        response.raise_for_status() # Will raise an HTTPError if the HTTP request returned an unsuccessful status code
        data = response.json()
        print("Response from image: ", data)
        # Expected structure: {'result': {'data': {'json': {'meta': {'prompt': '...', 'negativePrompt': '...'}}}}}
        meta = data.get('result', {}).get('data', {}).get('json', {}).get('meta')
        if meta: # meta can be None
            prompt = meta.get('prompt', "")
            negative_prompt = meta.get('negativePrompt', "")
    except requests.exceptions.RequestException as e:
        print(f"Could not fetch/parse generation data for image_id {image_id}: {e}")
    except json.JSONDecodeError as e:
        print(f"JSONDecodeError for image_id {image_id}: {e}. Response content: {response.text[:200]}")
    
    return prompt, negative_prompt

def extract_info(json_data: Dict[str, Any], hunyuan_type: Optional[str] = None) -> Optional[Dict[str, Any]]:
    if json_data.get("type") != "LORA":
        print("Model type is not LORA.")
        return None

    for model_version in json_data.get("modelVersions", []):
        civitai_base_model_name = model_version.get("baseModel")
        if civitai_base_model_name in SUPPORTED_CIVITAI_BASE_MODELS:
            base_model_hf = ""
            is_video = False

            if civitai_base_model_name == "Hunyuan Video":
                is_video = True
                if hunyuan_type == "Text-to-Video":
                    base_model_hf = "hunyuanvideo-community/HunyuanVideo"
                else: # Default or "Image-to-Video"
                    base_model_hf = "hunyuanvideo-community/HunyuanVideo-I2V"
            elif civitai_base_model_name in MODEL_MAPPING_VIDEO:
                is_video = True
                base_model_hf = MODEL_MAPPING_VIDEO[civitai_base_model_name]
            elif civitai_base_model_name in MODEL_MAPPING_IMAGE:
                base_model_hf = MODEL_MAPPING_IMAGE[civitai_base_model_name]
            else:
                print(f"Logic error: {civitai_base_model_name} in supported list but not mapped.")
                continue 

            primary_file_info = None
            for file_entry in model_version.get("files", []):
                if file_entry.get("primary", False) and file_entry.get("type") == "Model":
                    primary_file_info = file_entry
                    break
            
            if not primary_file_info:
                for file_entry in model_version.get("files", []):
                    if file_entry.get("type") == "Model" and file_entry.get("name","").endswith(".safetensors"):
                        primary_file_info = file_entry
                        print(f"Using first safetensors file as primary: {primary_file_info['name']}")
                        break
                if not primary_file_info:
                    print(f"No primary or suitable safetensors model file found for version {model_version.get('name')}")
                    continue

            urls_to_download = [{"url": primary_file_info["downloadUrl"], "filename": primary_file_info["name"], "type": "weightName"}]
            
            for image_obj in model_version.get("images", []):
                image_url = image_obj.get("url")
                if not image_url:
                    continue

                image_nsfw_level = image_obj.get("nsfwLevel", 0)
                if image_nsfw_level > 5:
                    continue
                    
                filename_part = os.path.basename(image_url)
                image_id_str = filename_part.split('.')[0]

                prompt, negative_prompt = "", ""
                if image_obj.get("hasMeta", False):
                     prompt, negative_prompt = get_prompts_from_image(image_id_str)

                urls_to_download.append({
                    "url": image_url,
                    "filename": filename_part, 
                    "type": "imageName", 
                    "prompt": prompt,
                    "negative_prompt": negative_prompt,
                    "media_type": image_obj.get("type", "image") 
                })

            info = {
                "urls_to_download": urls_to_download,
                "id": model_version["id"],
                "baseModel": base_model_hf, 
                "civitai_base_model_name": civitai_base_model_name,
                "is_video_model": is_video,
                "modelId": json_data.get("id", ""),
                "name": json_data["name"],
                "description": json_data.get("description", ""),
                "trainedWords": model_version.get("trainedWords", []),
                "creator": json_data.get("creator", {}).get("username", "Unknown"),
                "tags": json_data.get("tags", []),
                "allowNoCredit": json_data.get("allowNoCredit", True),
                "allowCommercialUse": json_data.get("allowCommercialUse", "Sell"),
                "allowDerivatives": json_data.get("allowDerivatives", True),
                "allowDifferentLicense": json_data.get("allowDifferentLicense", True)
            }
            return info
    print("No suitable model version found with a supported base model.")
    return None

def download_files(info, folder="."):
    downloaded_files = {
        "imageName": [], # Will contain both image and video filenames
        "imagePrompt": [],
        "imageNegativePrompt": [],
        "weightName": [],
        "mediaType": [] # To distinguish image/video for gallery if needed later
    }
    for item in info["urls_to_download"]:
        # Ensure filename is safe for filesystem
        safe_filename = slugify(item["filename"].rsplit('.', 1)[0]) + '.' + item["filename"].rsplit('.', 1)[-1] if '.' in item["filename"] else slugify(item["filename"])
        
        # Civitai URLs might need auth for direct download if not public
        try:
            download_file_with_auth(item["url"], safe_filename, folder) # Changed to use the auth-aware download
            downloaded_files[item["type"]].append(safe_filename)
            if item["type"] == "imageName": # This list now includes videos too
                prompt_clean = re.sub(r'<.*?>', '', item.get("prompt", ""))
                negative_prompt_clean = re.sub(r'<.*?>', '', item.get("negative_prompt", ""))
                downloaded_files["imagePrompt"].append(prompt_clean)
                downloaded_files["imageNegativePrompt"].append(negative_prompt_clean)
                downloaded_files["mediaType"].append(item.get("media_type", "image"))
        except gr.Error as e: # Catch Gradio errors from download_file_with_auth
            print(f"Skipping file {safe_filename} due to download error: {e.message}")
            gr.Warning(f"Skipping file {safe_filename} due to download error: {e.message}")

    return downloaded_files

# Renamed original download_file to download_file_with_auth
def download_file_with_auth(url, filename, folder="."):
    headers = {}
    # Add CIVITAI_API_TOKEN if available, for potentially restricted downloads
    # Note: The prompt example didn't use it for image URLs, only for the model file via API.
    # However, some image/video URLs might also require it if they are not fully public.
    if "CIVITAI_API_TOKEN" in os.environ: # Changed from CIVITAI_API
         headers['Authorization'] = f'Bearer {os.environ["CIVITAI_API_TOKEN"]}'

    try:
        response = requests.get(url, headers=headers, stream=True, timeout=60) # Added stream and timeout
        response.raise_for_status()
    except requests.exceptions.HTTPError as e:
        print(f"HTTPError downloading {url}: {e}")
        # No automatic retry with token here as it was specific to the primary file in original code
        # If it was related to auth, the initial header should have helped.
        raise gr.Error(f"Error downloading file {filename}: {e}")
    except requests.exceptions.RequestException as e:
        print(f"RequestException downloading {url}: {e}")
        raise gr.Error(f"Error downloading file {filename}: {e}")

    filepath = os.path.join(folder, filename)
    with open(filepath, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    print(f"Successfully downloaded {filepath}")


def process_url(url, profile, do_download=True, folder=".", hunyuan_type: Optional[str] = None):
    json_data = get_json_data(url)
    if json_data:
        # Always extract info, even if NSFW
        info = extract_info(json_data, hunyuan_type=hunyuan_type)
        if info:
            # Detect NSFW but do not block
            nsfw_flag = check_nsfw(json_data)
            info["nsfw_flag"] = nsfw_flag

            downloaded_files_summary = {}
            if do_download:
                gr.Info(f"Downloading files for {info['name']}...")
                downloaded_files_summary = download_files(info, folder)
                gr.Info(f"Finished downloading files for {info['name']}.")

            return info, downloaded_files_summary
        else:
            raise gr.Error("LoRA extraction failed. The base model might not be supported, or it's not a LoRA model, or no suitable files found in the version.")
    else:
        raise gr.Error("Failed to fetch model data from CivitAI API. Please check the URL and CivitAI's status.")


def create_readme(info: Dict[str, Any], downloaded_files: Dict[str, Any], user_repo_id: str, link_civit: bool = False, is_author: bool = True, folder: str = "."):
    readme_content = ""
    original_url = f"https://civitai.com/models/{info['modelId']}" if info.get('modelId') else "CivitAI (ID not found)"
    link_civit_disclaimer = f'([CivitAI]({original_url}))'
    non_author_disclaimer = f'This model was originally uploaded on [CivitAI]({original_url}), by [{info["creator"]}](https://civitai.com/user/{info["creator"]}/models). The information below was provided by the author on CivitAI:'
    
    is_video = info.get("is_video_model", False)
    base_hf_model = info["baseModel"] # This is the HF model ID
    civitai_bm_name_lower = info.get("civitai_base_model_name", "").lower()

    if is_video:
        default_tags = ["lora", "diffusers", "migrated", "video"]
        if "template:" not in " ".join(info.get("tags", [])):
             default_tags.append("template:video-lora")
        if "t2v" in civitai_bm_name_lower or (civitai_bm_name_lower == "hunyuan video" and base_hf_model.endswith("HunyuanVideo")):
            default_tags.append("text-to-video")
        elif "i2v" in civitai_bm_name_lower or (civitai_bm_name_lower == "hunyuan video" and base_hf_model.endswith("HunyuanVideo-I2V")):
            default_tags.append("image-to-video")
    else:
        default_tags = ["text-to-image", "stable-diffusion", "lora", "diffusers", "migrated"]
        if "template:" not in " ".join(info.get("tags", [])):
            default_tags.append("template:sd-lora")

    civit_tags_raw = info.get("tags", [])
    civit_tags_clean = [t.replace(":", "").strip() for t in civit_tags_raw if t.replace(":", "").strip()]
    final_civit_tags = [tag for tag in civit_tags_clean if tag not in default_tags and tag.lower() not in default_tags]
    tags = default_tags #+ final_civit_tags
    unpacked_tags = "\n- ".join(sorted(list(set(tags))))

    trained_words = info.get('trainedWords', [])
    formatted_words = ', '.join(f'`{word}`' for word in trained_words if word)
    trigger_words_section = f"## Trigger words\nYou should use {formatted_words} to trigger the generation." if formatted_words else ""
    
    widget_content = ""
    max_widget_items = 5
    items_for_widget = list(zip(
        downloaded_files.get("imagePrompt", []),
        downloaded_files.get("imageNegativePrompt", []),
        downloaded_files.get("imageName", [])
    ))[:max_widget_items]

    for index, (prompt, negative_prompt, media_filename) in enumerate(items_for_widget):
        escaped_prompt = prompt.replace("'", "''") if prompt else ' '
        base_media_filename = os.path.basename(media_filename)
        negative_prompt_content = f"negative_prompt: {negative_prompt}\n" if negative_prompt else ""
        # Corrected YAML for widget:
        widget_content += f"""- text: '{escaped_prompt}'
  {negative_prompt_content}
  output:
    url: >-
      {base_media_filename}
"""

    if base_hf_model in ["black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"]:
        dtype = "torch.bfloat16"
    else:
        dtype = "torch.float16" # Default for others, Hunyuan examples specify this.
    
    main_prompt_for_snippet_raw = formatted_words if formatted_words else 'Your custom prompt'
    if items_for_widget and items_for_widget[0][0]:
        main_prompt_for_snippet_raw = items_for_widget[0][0]
    
    # Escape single quotes for Python string literals
    main_prompt_for_snippet = main_prompt_for_snippet_raw.replace("'", "\\'")


    lora_loader_line = f"pipe.load_lora_weights('{user_repo_id}', weight_name='{downloaded_files.get('weightName', ['your_lora.safetensors'])[0]}')"

    diffusers_example = ""
    if is_video:
        if base_hf_model == "hunyuanvideo-community/HunyuanVideo-I2V":
            diffusers_example = f"""
```py
import torch
from diffusers import HunyuanVideoImageToVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import load_image, export_to_video

# Available checkpoints: "hunyuanvideo-community/HunyuanVideo-I2V" and "hunyuanvideo-community/HunyuanVideo-I2V-33ch"
model_id = "{base_hf_model}"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id, subfolder="transformer", torch_dtype=torch.bfloat16 # Explicitly bfloat16 for transformer
)
pipe = HunyuanVideoImageToVideoPipeline.from_pretrained(
    model_id, transformer=transformer, torch_dtype=torch.float16 # float16 for pipeline
)
pipe.vae.enable_tiling()
{lora_loader_line}
pipe.to("cuda")

prompt = "{main_prompt_for_snippet if main_prompt_for_snippet else 'A detailed scene description'}"
# Replace with your image path or URL
image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/guitar-man.png"
image = load_image(image_url)

output = pipe(image=image, prompt=prompt).frames[0]
export_to_video(output, "output.mp4", fps=15)
```
"""
        elif base_hf_model == "hunyuanvideo-community/HunyuanVideo":
            diffusers_example = f"""
```py
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

model_id = "{base_hf_model}"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id, subfolder="transformer", torch_dtype=torch.bfloat16
)
pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.float16)
{lora_loader_line}
# Enable memory savings
pipe.vae.enable_tiling()
pipe.enable_model_cpu_offload() # Optional: if VRAM is limited

output = pipe(
    prompt="{main_prompt_for_snippet if main_prompt_for_snippet else 'A cinematic video scene'}",
    height=320, # Adjust as needed
    width=512,  # Adjust as needed
    num_frames=61, # Adjust as needed
    num_inference_steps=30, # Adjust as needed
).frames[0]
export_to_video(output, "output.mp4", fps=15)
```
"""
        elif base_hf_model == "Lightricks/LTX-Video-0.9.7-dev" or base_hf_model == "Lightricks/LTX-Video-0.9.7-distilled": # Assuming -dev is the one from mapping
            # Note: The LTX example is complex. We'll simplify a bit for a LoRA example.
            # The user might need to adapt the full pipeline if they used the distilled one directly.
            # We assume the LoRA is trained on the main LTX pipeline.
            diffusers_example = f"""
```py
import torch
from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline
from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
from diffusers.utils import export_to_video, load_image, load_video

# Use the base LTX model your LoRA was trained on. The example below uses the distilled version.
# Adjust if your LoRA is for the non-distilled "Lightricks/LTX-Video-0.9.7-dev".
pipe = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.7-distilled", torch_dtype=torch.bfloat16)
{lora_loader_line}
# The LTX upsampler is separate and typically doesn't have LoRAs loaded into it directly.
pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained("Lightricks/ltxv-spatial-upscaler-0.9.7", vae=pipe.vae, torch_dtype=torch.bfloat16)

pipe.to("cuda")
pipe_upsample.to("cuda")
pipe.vae.enable_tiling()

def round_to_nearest_resolution_acceptable_by_vae(height, width, vae_spatial_compression_ratio):
    height = height - (height % vae_spatial_compression_ratio)
    width = width - (width % vae_spatial_compression_ratio)
    return height, width

# Example image for condition (replace with your own)
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png")
video_for_condition = load_video(export_to_video([image])) # Create a dummy video for conditioning
condition1 = LTXVideoCondition(video=video_for_condition, frame_index=0)

prompt = "{main_prompt_for_snippet if main_prompt_for_snippet else 'A cute little penguin takes out a book and starts reading it'}"
negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted" # Example
expected_height, expected_width = 480, 832 # Target final resolution
downscale_factor = 2 / 3
num_frames = 32 # Reduced for quicker example

# Part 1. Generate video at smaller resolution
downscaled_height, downscaled_width = int(expected_height * downscale_factor), int(expected_width * downscale_factor)
downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(downscaled_height, downscaled_width, pipe.vae_spatial_compression_ratio)

latents = pipe(
    conditions=[condition1],
    prompt=prompt,
    negative_prompt=negative_prompt,
    width=downscaled_width,
    height=downscaled_height,
    num_frames=num_frames,
    num_inference_steps=7, # Example steps
    guidance_scale=1.0,    # Example guidance
    decode_timestep = 0.05,
    decode_noise_scale = 0.025,
    generator=torch.Generator().manual_seed(0),
    output_type="latent",
).frames

# Part 2. Upscale generated video
upscaled_latents = pipe_upsample(
    latents=latents,
    output_type="latent"
).frames

# Part 3. Denoise the upscaled video (optional, but recommended)
video_frames = pipe(
    conditions=[condition1],
    prompt=prompt,
    negative_prompt=negative_prompt,
    width=downscaled_width * 2, # Upscaled width
    height=downscaled_height * 2, # Upscaled height
    num_frames=num_frames,
    denoise_strength=0.3,
    num_inference_steps=10,
    guidance_scale=1.0,
    latents=upscaled_latents,
    decode_timestep = 0.05,
    decode_noise_scale = 0.025,
    image_cond_noise_scale=0.025, # if using image condition
    generator=torch.Generator().manual_seed(0),
    output_type="pil",
).frames[0]

# Part 4. Downscale to target resolution if upscaler overshot
final_video = [frame.resize((expected_width, expected_height)) for frame in video_frames]
export_to_video(final_video, "output.mp4", fps=16) # Example fps
```
"""
        elif base_hf_model.startswith("Wan-AI/Wan2.1-T2V-"):
            diffusers_example = f"""
```py
import torch
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.utils import export_to_video

model_id = "{base_hf_model}"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) # As per example
pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
{lora_loader_line}
pipe.to("cuda")

prompt = "{main_prompt_for_snippet if main_prompt_for_snippet else 'A cat walks on the grass, realistic'}"
negative_prompt = "worst quality, low quality, blurry" # Simplified for LoRA example

output = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    height=480, # Adjust as needed
    width=832,  # Adjust as needed
    num_frames=30, # Adjust for LoRA, original example had 81
    guidance_scale=5.0 # Adjust as needed
).frames[0]
export_to_video(output, "output.mp4", fps=15)
```
"""
        elif base_hf_model.startswith("Wan-AI/Wan2.1-I2V-"):
            diffusers_example = f"""
```py
import torch
import numpy as np
from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
from diffusers.utils import export_to_video, load_image
from transformers import CLIPVisionModel

model_id = "{base_hf_model}"
# These components are part of the base model, LoRA is loaded into the pipeline
image_encoder = CLIPVisionModel.from_pretrained(model_id, subfolder="image_encoder", torch_dtype=torch.float32)
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
pipe = WanImageToVideoPipeline.from_pretrained(model_id, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16)
{lora_loader_line}
pipe.to("cuda")

# Replace with your image path or URL
image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
image = load_image(image_url)

# Adjust resolution based on model capabilities (480p or 720p variants)
# This is a simplified example; refer to original Wan I2V docs for precise resolution handling
if "480P" in model_id:
    max_height, max_width = 480, 832 # Example for 480p
elif "720P" in model_id:
    max_height, max_width = 720, 1280 # Example for 720p
else: # Fallback
    max_height, max_width = 480, 832

# Simple resize for example, optimal resizing might need to maintain aspect ratio & VAE constraints
h, w = image.height, image.width
if w > max_width or h > max_height:
    aspect_ratio = w / h
    if w > h:
        new_w = max_width
        new_h = int(new_w / aspect_ratio)
    else:
        new_h = max_height
        new_w = int(new_h * aspect_ratio)
    # Ensure dimensions are divisible by VAE scale factors (typically 8 or 16)
    # This is a basic adjustment, model specific patch sizes might also matter.
    patch_size_factor = 16 # Common factor
    new_h = (new_h // patch_size_factor) * patch_size_factor
    new_w = (new_w // patch_size_factor) * patch_size_factor
    if new_h > 0 and new_w > 0:
         image = image.resize((new_w, new_h))
    else: # Fallback if calculations lead to zero
        image = image.resize((max_width//2, max_height//2)) # A smaller safe default
else:
    patch_size_factor = 16 
    h = (h // patch_size_factor) * patch_size_factor
    w = (w // patch_size_factor) * patch_size_factor
    if h > 0 and w > 0:
        image = image.resize((w,h))


prompt = "{main_prompt_for_snippet if main_prompt_for_snippet else 'An astronaut in a dynamic scene'}"
negative_prompt = "worst quality, low quality, blurry" # Simplified

output = pipe(
    image=image, 
    prompt=prompt, 
    negative_prompt=negative_prompt, 
    height=image.height, # Use resized image height
    width=image.width,   # Use resized image width
    num_frames=30,       # Adjust for LoRA
    guidance_scale=5.0   # Adjust as needed
).frames[0]
export_to_video(output, "output.mp4", fps=16)
```
"""
        else: # Fallback for other video LoRAs
            diffusers_example = f"""
```py
# This is a video LoRA. Diffusers usage for video models can vary.
# You may need to install/import specific pipeline classes from diffusers or the model's community.
# Below is a generic placeholder.
import torch
from diffusers import AutoPipelineForTextToVideo # Or the appropriate video pipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

pipeline = AutoPipelineForTextToVideo.from_pretrained('{base_hf_model}', torch_dtype={dtype}).to(device)
{lora_loader_line}

# The following generation command is an example and may need adjustments
# based on the specific pipeline and its required parameters for '{base_hf_model}'.
# video_frames = pipeline(prompt='{main_prompt_for_snippet}', num_frames=16).frames
# For more details, consult the Hugging Face Hub page for {base_hf_model}
# and the Diffusers documentation on LoRAs and video pipelines.
```
"""
    else: # Image model
        diffusers_example = f"""
```py
from diffusers import AutoPipelineForText2Image
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"

pipeline = AutoPipelineForText2Image.from_pretrained('{base_hf_model}', torch_dtype={dtype}).to(device)
{lora_loader_line}
image = pipeline('{main_prompt_for_snippet}').images[0]
```
"""

    license_map_simple = {
        "Public Domain": "public-domain",
        "CreativeML Open RAIL-M": "creativeml-openrail-m",
        "CreativeML Open RAIL++-M": "creativeml-openrail-m", 
        "openrail": "creativeml-openrail-m",
    }
    commercial_use = info.get("allowCommercialUse", "None") 
    license_identifier = "other"
    license_name = "bespoke-lora-trained-license"
    
    if isinstance(commercial_use, str) and commercial_use.lower() == "none" and not info.get("allowDerivatives", True):
        license_identifier = "creativeml-openrail-m" 
        license_name = "CreativeML OpenRAIL-M" 
    
    bespoke_license_link = f"https://multimodal.art/civitai-licenses?allowNoCredit={info['allowNoCredit']}&allowCommercialUse={commercial_use[0] if isinstance(commercial_use, list) and commercial_use else (commercial_use if isinstance(commercial_use, str) else 'None')}&allowDerivatives={info['allowDerivatives']}&allowDifferentLicense={info['allowDifferentLicense']}"

    content = f"""---
license: {license_identifier}
license_name: "{license_name}"
license_link: {bespoke_license_link}
tags:
- {unpacked_tags}

base_model: {base_hf_model}
instance_prompt: '{trained_words[0] if trained_words else ''}'
widget:
{widget_content.strip()}
---

# {info["name"]} 

<Gallery />

{non_author_disclaimer if not is_author else ''}
{link_civit_disclaimer if link_civit else ''}

## Model description
{info["description"] if info["description"] else "No description provided."}

{trigger_words_section}

## Download model
Weights for this model are available in Safetensors format.
[Download](/{user_repo_id}/tree/main) them in the Files & versions tab.

## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers)
{diffusers_example}
For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters)
"""
    readme_content += content + "\n"
    readme_path = os.path.join(folder, "README.md")
    with open(readme_path, "w", encoding="utf-8") as file:
        file.write(readme_content)
    print(f"README.md created at {readme_path}")
    # print(f"README.md content:\n{readme_content}") # For debugging

def get_creator(username):
    url = f"https://civitai.com/api/trpc/user.getCreator?input=%7B%22json%22%3A%7B%22username%22%3A%22{username}%22%2C%22authed%22%3Atrue%7D%7D"
    try:
        response = requests.get(url, headers=headers, timeout=10)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error fetching creator data for {username}: {e}")
        gr.Warning(f"Could not verify Civitai creator's HF link: {e}")
        return None


def extract_huggingface_username(username_civitai):
    data = get_creator(username_civitai)
    if not data:
        return None
        
    links = data.get('result', {}).get('data', {}).get('json', {}).get('links', [])
    for link in links:
        url = link.get('url', '')
        if 'huggingface.co/' in url:
            # Extract username, handling potential variations like www. or trailing slashes
            hf_username = url.split('huggingface.co/')[-1].split('/')[0]
            if hf_username:
                return hf_username
    return None


def check_civit_link(profile: Optional[gr.OAuthProfile], url: str):
    # Initial return structure: instructions_html, submit_interactive, try_again_visible, other_submit_visible, hunyuan_radio_visible
    # Default to disabling/hiding things if checks fail early
    default_fail_updates = ("", gr.update(interactive=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False))

    if not profile: # Should be handled by demo.load and login button
        return "Please log in with Hugging Face.", gr.update(interactive=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

    if not url or not url.startswith("https://civitai.com/models/"):
        return "Please enter a valid Civitai model URL.", gr.update(interactive=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

    try:
        # We need hunyuan_type for extract_info, but we don't know it yet.
        # Call get_json_data first to check if it's Hunyuan.
        json_data_preview = get_json_data(url)
        if not json_data_preview:
            return ("Failed to fetch basic model info from Civitai. Check URL.",
                    gr.update(interactive=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False))

        is_hunyuan = False
        original_civitai_base_model = ""
        if json_data_preview.get("type") == "LORA":
            for mv in json_data_preview.get("modelVersions", []):
                # Try to find a relevant model version to check its base model
                # This is a simplified check; extract_info does a more thorough search
                cbm = mv.get("baseModel")
                if cbm and cbm in SUPPORTED_CIVITAI_BASE_MODELS:
                    original_civitai_base_model = cbm
                    if cbm == "Hunyuan Video":
                        is_hunyuan = True
                    break 
        
        # Now call process_url with a default hunyuan_type for other checks
        # The actual hunyuan_type choice will be used during the main upload.
        info, _ = process_url(url, profile, do_download=False, hunyuan_type="Image-to-Video") # Use default for check
        
        # If process_url raises an error (e.g. NSFW, not supported), it will be caught by Gradio
        # and displayed as a gr.Error. Here, we assume it passed if no exception.

    except gr.Error as e: # Catch errors from process_url (like NSFW, not supported)
        return (f"Cannot process this model: {e.message}",
                gr.update(interactive=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=is_hunyuan)) # Show hunyuan if detected
    except Exception as e: # Catch any other unexpected error during preview
        print(f"Unexpected error in check_civit_link: {e}")
        return (f"An unexpected error occurred: {str(e)}",
                gr.update(interactive=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=is_hunyuan))


    hf_username_on_civitai = extract_huggingface_username(info['creator'])
    
    if profile.username in TRUSTED_UPLOADERS:
        return ('Admin/Trusted user override: Upload enabled.', 
                gr.update(interactive=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=is_hunyuan))
        
    if not hf_username_on_civitai:
        no_username_text = (f'If you are {info["creator"]} on Civitai, hi! Your CivitAI profile does not seem to have a link to your Hugging Face account. '
                            f'Please visit <a href="https://civitai.com/user/account" target="_blank">https://civitai.com/user/account</a>, '
                            f'go to "Edit profile" and add your Hugging Face profile URL (e.g., https://huggingface.co/{profile.username}) to the "Links" section. '
                            f'<br><img width="60%" src="https://i.imgur.com/hCbo9uL.png" alt="Civitai profile links example"/><br>'
                            f'(If you are not {info["creator"]}, you cannot submit their model at this time.)')
        return no_username_text, gr.update(interactive=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=is_hunyuan)

    if profile.username.lower() != hf_username_on_civitai.lower():
        unmatched_username_text = (f'Oops! The Hugging Face username found on the CivitAI profile of {info["creator"]} is '
                                   f'"{hf_username_on_civitai}", but you are logged in as "{profile.username}". '
                                   f'Please ensure your CivitAI profile links to the correct Hugging Face account: '
                                   f'<a href="https://civitai.com/user/account" target="_blank">https://civitai.com/user/account</a> (Edit profile -> Links section).'
                                   f'<br><img width="60%" src="https://i.imgur.com/hCbo9uL.png" alt="Civitai profile links example"/>')
        return unmatched_username_text, gr.update(interactive=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=is_hunyuan)
    
    # All checks passed
    return ('Username verified! You can now upload this model.', 
            gr.update(interactive=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=is_hunyuan))

        
def swap_fill(profile: Optional[gr.OAuthProfile]):
    if profile is None: # Not logged in
        return gr.update(visible=True), gr.update(visible=False)
    else: # Logged in
        return gr.update(visible=False), gr.update(visible=True)

def show_output():
    return gr.update(visible=True)

def list_civit_models(username_civitai: str):
    if not username_civitai:
        return ""
    url = f"https://civitai.com/api/v1/models?username={username_civitai}&limit=100&sort=Newest" # Added sort
    
    all_model_urls = ""
    page_count = 0
    max_pages = 5 # Limit number of pages to fetch to avoid very long requests

    while url and page_count < max_pages:
        try:
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            data = response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error fetching model list for {username_civitai}: {e}")
            gr.Warning(f"Could not fetch full model list for {username_civitai}.")
            break 
        
        items = data.get('items', [])
        if not items:
            break

        for model in items:
            # Only list LORAs of supported base model types to avoid cluttering with unsupported ones
            is_supported_lora = False
            if model.get("type") == "LORA":
                # Check modelVersions for baseModel compatibility
                for mv in model.get("modelVersions", []):
                    if mv.get("baseModel") in SUPPORTED_CIVITAI_BASE_MODELS:
                        is_supported_lora = True
                        break
            if is_supported_lora:
                model_slug = slugify(model.get("name", f"model-{model['id']}"))
                all_model_urls += f'https://civitai.com/models/{model["id"]}/{model_slug}\n'
        
        metadata = data.get('metadata', {})
        url = metadata.get('nextPage', None)
        page_count += 1
        if page_count >= max_pages and url:
            print(f"Reached max page limit for fetching models for {username_civitai}.")
            gr.Info(f"Showing first {max_pages*100} models. There might be more.")

    if not all_model_urls:
        gr.Info(f"No compatible LoRA models found for user {username_civitai} or user not found.")
    return all_model_urls.strip()


def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: Optional[gr.OAuthToken], url: str, link_civit: bool, hunyuan_type: str):
    if not profile or not profile.username: # Check profile and username
        raise gr.Error("You must be logged in to Hugging Face to upload.")
    if not oauth_token or not oauth_token.token:
        raise gr.Error("Hugging Face authentication token is missing or invalid. Please log out and log back in.")
    
    folder = str(uuid.uuid4())
    os.makedirs(folder, exist_ok=True) # exist_ok=True is safer if folder might exist
    
    gr.Info(f"Starting processing for model {url}")
    try:
        # Pass hunyuan_type to process_url
        info, downloaded_files_summary = process_url(url, profile, do_download=True, folder=folder, hunyuan_type=hunyuan_type)
    except gr.Error as e: # Catch errors from process_url (NSFW, not supported, API fail)
        # Cleanup created folder if download failed or was skipped
        if os.path.exists(folder):
            try:
                import shutil
                shutil.rmtree(folder)
            except Exception as clean_e:
                print(f"Error cleaning up folder {folder}: {clean_e}")
        raise e # Re-raise the Gradio error to display it

    if not downloaded_files_summary.get("weightName"):
        raise gr.Error("No model weight file was downloaded. Cannot proceed with upload.")

    # Determine if user is the author for README generation
    # This relies on extract_huggingface_username which needs COOKIE_INFO
    is_author = False
    if "COOKIE_INFO" in os.environ:
        hf_username_on_civitai = extract_huggingface_username(info['creator'])
        if hf_username_on_civitai and profile.username.lower() == hf_username_on_civitai.lower():
            is_author = True
    elif profile.username.lower() == info['creator'].lower(): # Fallback if cookie not set, direct match
        is_author = True


    slug_name = slugify(info["name"])
    user_repo_id = f"{profile.username}/{slug_name}"
    
    gr.Info(f"Creating README for {user_repo_id}...")
    create_readme(info, downloaded_files_summary, user_repo_id, link_civit, is_author, folder=folder)
    
    try:
        gr.Info(f"Creating repository {user_repo_id} on Hugging Face...")
        create_repo(repo_id=user_repo_id, private=True, exist_ok=True, token=oauth_token.token)
        
        gr.Info(f"Starting upload of all files to {user_repo_id}...")
        upload_folder(
            folder_path=folder,
            repo_id=user_repo_id,
            repo_type="model",
            token=oauth_token.token,
            commit_message=f"Upload LoRA: {info['name']} from Civitai model ID {info['modelId']}" # Add commit message
        )
        
        gr.Info(f"Setting repository {user_repo_id} to public...")
        update_repo_visibility(repo_id=user_repo_id, private=False, token=oauth_token.token)
        gr.Info(f"Model {info['name']} uploaded successfully to {user_repo_id}!")
    except Exception as e:
        print(f"Error during Hugging Face repo operations for {user_repo_id}: {e}")
        # Attempt to provide a more specific error message for token issues
        if "401" in str(e) or "Unauthorized" in str(e):
             raise gr.Error("Hugging Face authentication failed (e.g. token expired or insufficient permissions). Please log out and log back in with a token that has write permissions.")
        raise gr.Error(f"Error during Hugging Face upload: {str(e)}")
    finally:
        # Clean up the temporary folder
        if os.path.exists(folder):
            try:
                import shutil
                shutil.rmtree(folder)
                print(f"Cleaned up temporary folder: {folder}")
            except Exception as clean_e:
                print(f"Error cleaning up folder {folder}: {clean_e}")
        
    return f"""# Model uploaded to 🤗!
Access it here: [{user_repo_id}](https://huggingface.co/{user_repo_id})
"""

def bulk_upload(profile: Optional[gr.OAuthProfile], oauth_token: Optional[gr.OAuthToken], urls_text: str, link_civit: bool, hunyuan_type: str):
    if not urls_text.strip():
        return "No URLs provided for bulk upload."
        
    urls = [url.strip() for url in urls_text.split("\n") if url.strip()]
    if not urls:
        return "No valid URLs found in the input."

    upload_results_md = "## Bulk Upload Results:\n\n"
    success_count = 0
    failure_count = 0

    for i, url in enumerate(urls):
        gr.Info(f"Processing URL {i+1}/{len(urls)}: {url}")
        try:
            result = upload_civit_to_hf(profile, oauth_token, url, link_civit, hunyuan_type)
            upload_results_md += f"**SUCCESS**: {url}\n{result}\n\n---\n\n"
            success_count +=1
        except gr.Error as e: # Catch Gradio-raised errors (expected failures)
            upload_results_md += f"**FAILED**: {url}\n*Reason*: {e.message}\n\n---\n\n"
            gr.Warning(f"Failed to upload {url}: {e.message}")
            failure_count +=1
        except Exception as e: # Catch unexpected Python errors
            upload_results_md += f"**FAILED**: {url}\n*Unexpected Error*: {str(e)}\n\n---\n\n"
            gr.Warning(f"Unexpected error uploading {url}: {str(e)}")
            failure_count +=1
            
    summary = f"Finished bulk upload: {success_count} successful, {failure_count} failed."
    gr.Info(summary)
    upload_results_md = f"## {summary}\n\n" + upload_results_md
    return upload_results_md

# --- Gradio UI ---
css = '''
#login_button_row button { /* Target login button specifically */
    width: 100% !important;
    margin: 0 auto;
}
#disabled_upload_area { /* ID for the disabled area */
    opacity: 0.5;
    pointer-events: none;
}
'''

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo: # Added a theme
    gr.Markdown('''# Upload your CivitAI LoRA to Hugging Face 🤗
By uploading your LoRAs to Hugging Face you get diffusers compatibility, a free GPU-based Inference Widget (for many models)
''')
    
    with gr.Row(elem_id="login_button_row"):
        login_button = gr.LoginButton() # Moved login_button definition here

    # Area shown when not logged in (or login fails)
    with gr.Column(elem_id="disabled_upload_area", visible=True) as disabled_area:
        gr.HTML("<i>Please log in with Hugging Face to enable uploads.</i>")
        # Add some dummy placeholders to mirror the enabled_area structure if needed for consistent layout
        gr.Textbox(label="CivitAI model URL (Log in to enable)", interactive=False)
        gr.Button("Upload (Log in to enable)", interactive=False)

    # Area shown when logged in
    with gr.Column(visible=False) as enabled_area:
        with gr.Row():
            submit_source_civit_enabled = gr.Textbox(
                placeholder="https://civitai.com/models/144684/pixelartredmond-pixel-art-loras-for-sd-xl",
                label="CivitAI model URL",
                info="URL of the CivitAI LoRA model page.",
                elem_id="submit_source_civit_main" # Unique ID
            )
        
        hunyuan_type_radio = gr.Radio(
            choices=["Image-to-Video", "Text-to-Video"],
            label="HunyuanVideo Type (Select if model is Hunyuan Video)",
            value="Image-to-Video", # Default as per prompt
            visible=False, # Initially hidden
            interactive=True
        )
        
        link_civit_checkbox = gr.Checkbox(label="Link back to original CivitAI page in README?", value=False)

        with gr.Accordion("Bulk Upload (Multiple LoRAs)", open=False):
            civit_username_to_bulk = gr.Textbox(
                label="Your CivitAI Username (Optional)",
                info="Type your CivitAI username here to automatically populate the list below with your compatible LoRAs."
            )
            submit_bulk_civit_urls = gr.Textbox(
                label="CivitAI Model URLs (One per line)",
                info="Add one CivitAI model URL per line for bulk processing.",
                lines=6,
            )
            bulk_button = gr.Button("Start Bulk Upload")
                
        instructions_html = gr.HTML("") # For messages from check_civit_link
        
        # Buttons for single upload
        # try_again_button is shown if username check fails
        try_again_button_single = gr.Button("I've updated my CivitAI profile, check again", visible=False)
        # submit_button_single is the main upload button for single model
        submit_button_single = gr.Button("Upload Model to Hugging Face", interactive=False, variant="primary")
        
        output_markdown = gr.Markdown(label="Upload Progress & Results", visible=False)

    # Event Handling
    # When login status changes (login_button implicitly handles profile state for demo.load)
    # demo.load updates visibility of disabled_area and enabled_area based on login.
    # The `profile` argument is implicitly passed by Gradio to functions that declare it.
    # `oauth_token` is also implicitly passed if `login_button` is used and function expects `gr.OAuthToken`.

    # When URL changes in the enabled area
    submit_source_civit_enabled.change(
        fn=check_civit_link,
        inputs=[submit_source_civit_enabled], # profile is implicitly passed
        outputs=[instructions_html, submit_button_single, try_again_button_single, submit_button_single, hunyuan_type_radio],
        # Outputs map to: instructions, submit_interactive, try_again_visible, (submit_visible - seems redundant here, check_civit_link logic ensures one is visible), hunyuan_radio_visible
        # For submit_button_single: 2nd output controls 'interactive', 4th controls 'visible' (often paired with try_again_button's visibility)
    )

    # Try again button for single upload (re-checks the same URL)
    try_again_button_single.click(
        fn=check_civit_link,
        inputs=[submit_source_civit_enabled],
        outputs=[instructions_html, submit_button_single, try_again_button_single, submit_button_single, hunyuan_type_radio],
    )

    # Autofill bulk URLs from CivitAI username
    civit_username_to_bulk.change(
        fn=list_civit_models,
        inputs=[civit_username_to_bulk],
        outputs=[submit_bulk_civit_urls]
    )

    # Single model upload button click
    submit_button_single.click(fn=show_output, outputs=[output_markdown]).then(
        fn=upload_civit_to_hf,
        inputs=[submit_source_civit_enabled, link_civit_checkbox, hunyuan_type_radio], # profile, oauth_token implicit
        outputs=[output_markdown]
    )

    # Bulk model upload button click
    bulk_button.click(fn=show_output, outputs=[output_markdown]).then(
        fn=bulk_upload,
        inputs=[submit_bulk_civit_urls, link_civit_checkbox, hunyuan_type_radio], # profile, oauth_token implicit
        outputs=[output_markdown]
    )
    
    # Initial state of visible areas based on login status
    demo.load(fn=swap_fill, outputs=[disabled_area, enabled_area], queue=False)

demo.queue(default_concurrency_limit=5) # Reduced concurrency from 50, can be demanding
demo.launch(debug=True) # Added debug=True for development