Update app.py

app.py

@@ -314,39 +314,34 @@ def create_readme(info: Dict[str, Any], downloaded_files: Dict[str, Any], user_r
     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"]
+    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["tags"]): # if no template tag from civitai
-             default_tags.append("template:video-lora") # A generic video template tag
+        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["tags"]):
+        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()] # Clean and remove empty
-    # Filter out tags already covered by default_tags logic (e.g. 'text-to-image', 'lora')
+    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)))) # Sort and unique
+    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) # Filter out empty/None words
+    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 = ""
-    # 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", []),
@@ -355,70 +350,292 @@ def create_readme(info: Dict[str, Any], downloaded_files: Dict[str, Any], user_r
     ))[: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
+        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}  
+  {negative_prompt}
   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"
+        dtype = "torch.float16" # Default for others, Hunyuan examples specify this.
     
-    # 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]
+    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]
     
-    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
+    # Escape single quotes for Python string literals
+    main_prompt_for_snippet = main_prompt_for_snippet_raw.replace("'", "\\'")
 
 
-        diffusers_example = f"""
+    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.
-# 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
+# You may need to install/import specific pipeline classes from diffusers or the model's community.
+# Below is a generic placeholder.
 import torch
-# For Hunyuan, you might need: from hunyuanvideo_community.pipelines import {pipeline_class}
+from diffusers import AutoPipelineForTextToVideo # Or the appropriate video pipeline
 
 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]}')
+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.
-# {output_name} = pipeline({example_input}){output_access}
+# 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.
 ```
@@ -432,7 +649,7 @@ 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]}')
+{lora_loader_line}
 image = pipeline('{main_prompt_for_snippet}').images[0]
 ```
 """
@@ -440,31 +657,22 @@ 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
+        "CreativeML Open RAIL++-M": "creativeml-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
+    commercial_use = info.get("allowCommercialUse", "None") 
     license_identifier = "other"
-    license_name = "bespoke-lora-trained-license" # Default bespoke license
+    license_name = "bespoke-lora-trained-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_identifier = "creativeml-openrail-m" 
         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_name: "{license_name}"
 license_link: {bespoke_license_link}
 tags:
 - {unpacked_tags}
@@ -472,7 +680,8 @@ tags:
 base_model: {base_hf_model}
 instance_prompt: {trained_words[0] if trained_words else ''}
 widget:
-{widget_content}---
+{widget_content.strip()}
+---
 
 # {info["name"]} 
 
@@ -496,11 +705,10 @@ For more details, including weighting, merging and fusing LoRAs, check the [docu
 """
     readme_content += content + "\n"
     readme_path = os.path.join(folder, "README.md")
-    with open(readme_path, "w", encoding="utf-8") as file: # Added encoding
+    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 {readme_content}")
-
+    # 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"