Update app.py

app.py

@@ -42,6 +42,21 @@ MODEL_MAPPING_VIDEO = {
 
 SUPPORTED_CIVITAI_BASE_MODELS = list(MODEL_MAPPING_IMAGE.keys()) + list(MODEL_MAPPING_VIDEO.keys())
 
+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
+}
 
 def get_json_data(url):
     url_split = url.split('/')
@@ -112,9 +127,10 @@ def get_prompts_from_image(image_id_str: str):
     prompt = ""
     negative_prompt = ""
     try:
-        response = requests.get(url, timeout=10) # Added timeout
+        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
@@ -150,7 +166,6 @@ def extract_info(json_data: Dict[str, Any], hunyuan_type: Optional[str] = None)
             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 
 
@@ -161,7 +176,6 @@ def extract_info(json_data: Dict[str, Any], hunyuan_type: Optional[str] = None)
                     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
@@ -171,54 +185,43 @@ def extract_info(json_data: Dict[str, Any], hunyuan_type: Optional[str] = None)
                     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
+                filename_part = os.path.basename(image_url)
                 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
+                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, # Use the extracted filename part
-                    "type": "imageName", # Keep as imageName for consistency with README generation
+                    "filename": filename_part, 
+                    "type": "imageName", 
                     "prompt": prompt,
                     "negative_prompt": negative_prompt,
-                    "media_type": image_obj.get("type", "image") # store if it's 'image' or 'video'
+                    "media_type": image_obj.get("type", "image") 
                 })
 
             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
+                "baseModel": base_model_hf, 
+                "civitai_base_model_name": civitai_base_model_name,
                 "is_video_model": is_video,
-                "modelId": json_data.get("id", ""), # Main model ID from Civitai
+                "modelId": json_data.get("id", ""),
                 "name": json_data["name"],
-                "description": json_data.get("description", ""), # Description can be None
+                "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"), # Default to most permissive if missing
+                "allowCommercialUse": json_data.get("allowCommercialUse", "Sell"),
                 "allowDerivatives": json_data.get("allowDerivatives", True),
                 "allowDifferentLicense": json_data.get("allowDifferentLicense", True)
             }
@@ -507,21 +510,6 @@ def get_creator(username):
         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()