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 = ["KappaNeuro", "CiroN2022", "multimodalart", "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",
}

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())


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], profile: Optional[gr.OAuthProfile]) -> bool:
    if not json_data:
        return False # Should not happen if get_json_data succeeded

    # Overall model boolean flag - highest priority
    if json_data.get("nsfw", False):
        print("Model flagged as NSFW by 'nsfw: true'.")
        gr.Info("Reason: Model explicitly flagged as NSFW on Civitai.")
        return False # Unsafe

    # Overall model numeric nsfwLevel - second priority. Max allowed is 5 (nsfwLevel < 6).
    # nsfwLevel definitions: None (1), Mild (2), Mature (4), Adult (5), X (8), R (16), XXX (32)
    model_nsfw_level = json_data.get("nsfwLevel", 0)
    if model_nsfw_level > 5: # Anything above "Adult"
        print(f"Model's overall nsfwLevel ({model_nsfw_level}) is > 5. Blocking.")
        gr.Info(f"Reason: Model's overall NSFW Level ({model_nsfw_level}) is above the allowed threshold (5).")
        return False # Unsafe

    # If uploader is trusted and the above checks passed, they bypass further version/image checks.
    if profile and profile.username in TRUSTED_UPLOADERS:
        print(f"User {profile.username} is trusted. Model 'nsfw' is false and overall nsfwLevel ({model_nsfw_level}) is <= 5. Allowing.")
        return True

    # For non-trusted users, check nsfwLevel of model versions and individual images/videos
    for model_version in json_data.get("modelVersions", []):
        version_nsfw_level = model_version.get("nsfwLevel", 0)
        if version_nsfw_level > 5:
            print(f"Model version nsfwLevel ({version_nsfw_level}) is > 5 for non-trusted user. Blocking.")
            gr.Info(f"Reason: A model version's NSFW Level ({version_nsfw_level}) is above 5.")
            return False
        for image_item in model_version.get("images", []):
            item_nsfw_level = image_item.get("nsfwLevel", 0)
            if item_nsfw_level > 5:
                print(f"Media item nsfwLevel ({item_nsfw_level}) is > 5 for non-trusted user. Blocking.")
                gr.Info(f"Reason: An example image/video's NSFW Level ({item_nsfw_level}) is above 5.")
                return False
    return True # Safe for non-trusted user if all checks pass


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, timeout=10) # Added timeout
        response.raise_for_status() # Will raise an HTTPError if the HTTP request returned an unsuccessful status code
        data = response.json()
        # 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:
                # Should not happen if SUPPORTED_CIVITAI_BASE_MODELS is derived correctly
                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:
                # Sometimes primary might not be explicitly set, take first 'Model' type safetensors
                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", []):
                # Skip if image/video itself is too NSFW for non-trusted, extract_info is called by check_civit_link too early for nsfw check
                # The main nsfw check will handle this before download. Here we just gather info.
                # if image_obj.get("nsfwLevel", 0) > 5: # This check belongs in check_nsfw for non-trusted.
                #     continue
                
                image_url = image_obj.get("url")
                if not image_url:
                    continue

                # Extract image ID for fetching prompts. This can be fragile.
                # Example URLs:
                # https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/.../12345.jpeg (where 12345 is the id)
                # https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/.../width=1024/12345.jpeg
                # https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/.../12345.mp4
                filename_part = os.path.basename(image_url) # e.g., 12345.jpeg or 12345.mp4
                image_id_str = filename_part.split('.')[0]

                prompt, negative_prompt = "", ""
                if image_obj.get("hasMeta", False) and image_obj.get("type") == "image": # Often only images have reliable meta via this endpoint
                     prompt, negative_prompt = get_prompts_from_image(image_id_str)

                urls_to_download.append({
                    "url": image_url,
                    "filename": filename_part, # Use the extracted filename part
                    "type": "imageName", # Keep as imageName for consistency with README generation
                    "prompt": prompt,
                    "negative_prompt": negative_prompt,
                    "media_type": image_obj.get("type", "image") # store if it's 'image' or 'video'
                })

            info = {
                "urls_to_download": urls_to_download,
                "id": model_version["id"],
                "baseModel": base_model_hf, # This is the HF model ID
                "civitai_base_model_name": civitai_base_model_name, # Original name from Civitai
                "is_video_model": is_video,
                "modelId": json_data.get("id", ""), # Main model ID from Civitai
                "name": json_data["name"],
                "description": json_data.get("description", ""), # Description can be None
                "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"), # Default to most permissive if missing
                "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:
        if check_nsfw(json_data, profile):
            info = extract_info(json_data, hunyuan_type=hunyuan_type)
            if info:
                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:
            # check_nsfw now prints detailed reasons via gr.Info/print
            raise gr.Error("This model has content tagged as unsafe by CivitAI or exceeds NSFW level limits.")
    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:'
    
    # Tags
    is_video = info.get("is_video_model", False)
    base_hf_model = info["baseModel"]
    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["tags"]): # if no template tag from civitai
             default_tags.append("template:video-lora") # A generic video template tag
        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["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()] # Clean and remove empty
    # Filter out tags already covered by default_tags logic (e.g. 'text-to-image', 'lora')
    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)))) # Sort and unique

    trained_words = info.get('trainedWords', [])
    formatted_words = ', '.join(f'`{word}`' for word in trained_words if word) # Filter out empty/None words
    trigger_words_section = f"## Trigger words\nYou should use {formatted_words} to trigger the generation." if formatted_words else ""
    
    widget_content = ""
    # Limit number of widget items to avoid overly long READMEs, e.g., max 5
    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 ' ' # Handle None or empty prompt
        
        # Ensure media_filename is just the filename, not a path
        base_media_filename = os.path.basename(media_filename)

        negative_prompt_content = f"    negative_prompt: {negative_prompt}\n" if negative_prompt else ""
        widget_content += f"""- text: '{escaped_prompt}'
  {negative_prompt_content}  
  output:
    url: >-
      {base_media_filename}
"""
    # Determine dtype
    if base_hf_model in ["black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"]:
        dtype = "torch.bfloat16"
    else:
        dtype = "torch.float16"
    
    # Diffusers code snippet
    main_prompt_for_snippet = formatted_words if formatted_words else 'Your custom prompt'
    # If a specific prompt exists for the first image/video, use that one
    if items_for_widget and items_for_widget[0][0]: # items_for_widget[0][0] is the prompt of the first media
        main_prompt_for_snippet = items_for_widget[0][0]
    
    if is_video:
        # Determine if T2V or I2V for example snippet based on HF model name or Civitai name
        pipeline_class = "AutoPipelineForTextToVideo" # Default for T2V
        example_input = f"'{main_prompt_for_snippet}'"
        output_name = "video_frames"
        output_access = ".frames"

        if "I2V" in base_hf_model or "i2v" in civitai_bm_name_lower:
            pipeline_class = "AutoPipelineForVideoToVideo" # Or ImageToVideo if more specific class exists
            example_input = f"prompt='{main_prompt_for_snippet}', image=your_input_image_or_pil" # I2V needs an image
            # For I2V, .frames might still be correct but input changes.
        
        # Handle Hunyuan specifically for more accurate snippet if possible
        if "HunyuanVideo" in base_hf_model:
            if base_hf_model.endswith("HunyuanVideo"): # T2V
                pipeline_class = "HunyuanDiT2V Pipeline" # from hunyuanvideo_community.pipelines.hunyuan_dit_t2v_pipeline import HunyuanDiT2V Pipeline
                example_input = f"prompt='{main_prompt_for_snippet}', height=576, width=1024, num_frames=16, num_inference_steps=50, guidance_scale=7.5" # Example params
            else: # I2V
                pipeline_class = "HunyuanDiI2V Pipeline" # from hunyuanvideo_community.pipelines.hunyuan_dit_i2v_pipeline import HunyuanDiI2V Pipeline
                example_input = f"pil_image, prompt='{main_prompt_for_snippet}', height=576, width=1024, num_frames=16, num_inference_steps=50, guidance_scale=7.5, strength=0.8" # Example params


        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.
# Example for a {pipeline_class.split()[0]} based workflow:
from diffusers import {pipeline_class.split()[0]} # Adjust if pipeline_class includes more than just class name
import torch
# For Hunyuan, you might need: from hunyuanvideo_community.pipelines import {pipeline_class}

device = "cuda" if torch.cuda.is_available() else "cpu"
# pil_image = ... # Load your input image PIL here if it's an Image-to-Video model

pipeline = {pipeline_class.split()[0]}.from_pretrained('{base_hf_model}', torch_dtype={dtype}).to(device)
pipeline.load_lora_weights('{user_repo_id}', weight_name='{downloaded_files["weightName"][0]}')

# The following generation command is an example and may need adjustments
# based on the specific pipeline and its required parameters.
# {output_name} = pipeline({example_input}){output_access}
# 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)
pipeline.load_lora_weights('{user_repo_id}', weight_name='{downloaded_files["weightName"][0]}')
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", # Assuming mapping to openrail-m
        "openrail": "creativeml-openrail-m",
        "SDXL": "sdxl", # This might be a base model, not a license
        # Add more mappings if CivitAI provides other common license names
    }
    # Attempt to map commercial use if possible, otherwise use bespoke
    # "allowCommercialUse": ["Image", "RentCivit", "Rent", "Sell"] or "None", "Sell" etc.
    commercial_use = info.get("allowCommercialUse", "None") # Default to None if not specified
    license_identifier = "other"
    license_name = "bespoke-lora-trained-license" # Default bespoke license
    
    # Heuristic for common licenses based on permissions
    if isinstance(commercial_use, str) and commercial_use.lower() == "none" and not info.get("allowDerivatives", True):
        license_identifier = "creativeml-openrail-m" # Or a more restrictive one if known
        license_name = "CreativeML OpenRAIL-M" 
    elif isinstance(commercial_use, list) and "Sell" in commercial_use and info.get("allowDerivatives", True):
        # This is a very permissive license, could be Apache 2.0 or MIT if source code, but for models, 'other' is safer
        pass # Keep bespoke for now

    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}---

# {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: # Added encoding
        file.write(readme_content)
    print(f"README.md created at {readme_path}")
    print(f"README.md content {readme_content}")


def get_creator(username):
    # Ensure COOKIE_INFO is set as an environment variable
    # Example: "__Host-next-auth.csrf-token=xxx; __Secure-next-auth.callback-url=yyy; __Secure-next-auth.session-token=zzz"
    cookie_info = os.environ.get("COOKIE_INFO")
    if not cookie_info:
        print("COOKIE_INFO environment variable not set. Cannot fetch creator's HF username.")
        gr.Warning("COOKIE_INFO not set. Cannot verify Hugging Face username on Civitai.")
        return None # Cannot proceed without cookie for this specific call

    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"
    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
        "referer": f"https://civitai.com/user/{username}/models",
        "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
    }
    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 == "multimodalart" or profile.username in TRUSTED_UPLOADERS: # Allow multimodalart or other trusted to bypass HF username check
        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 unless you are a trusted uploader.)')
        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