Update app.py

I DONT EVEN KNOW IF THIS WORKS BUT THIS IS THE OLD VERSION BUT YET NEW SINCE THE IDE KEPT EATING MY FILES

app.py

@@ -1,16 +1,10 @@
-# ---------------------- IMPORTS ----------------------
-# Core functionality
 import os
 import gradio as gr
 import torch
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL
 from transformers import CLIPTextModel
-
-# Model handling
 from safetensors.torch import load_file
 from collections import OrderedDict
-
-# Utilities
 import re
 import json
 import gdown
@@ -26,142 +20,153 @@ import shutil
 import hashlib
 from datetime import datetime
 from typing import Dict, List, Optional
+from huggingface_hub import login, HfApi, validate_repo_id, HfHubHTTPError
 
-# Hugging Face integration
-from huggingface_hub import login, HfApi
-from types import SimpleNamespace
+# ---------------------- DEPENDENCIES ----------------------
+def install_dependencies_gradio():
+    """Installs the necessary dependencies for the Gradio app. Run this ONCE."""
+    try:
+        subprocess.run(["pip", "install", "-U", "torch", "diffusers", "transformers", "accelerate", "safetensors", "huggingface_hub", "xformers"])
+        print("Dependencies installed successfully.")
+    except Exception as e:
+        print(f"Error installing dependencies: {e}")
 
 # ---------------------- UTILITY FUNCTIONS ----------------------
-def is_valid_url(url):
-    """Check if a string is a valid URL."""
-    try:
-        result = urlparse(url)
-        return all([result.scheme, result.netloc])
-    except:
-        return False
+def get_save_dtype(save_precision_as):
+    """Determines the save dtype based on the user's choice."""
+    if save_precision_as == "fp16":
+        return torch.float16
+    elif save_precision_as == "bf16":
+        return torch.bfloat16
+    elif save_precision_as == "float":
+        return torch.float32
+    else:
+        return None
 
-def get_filename(url):
-    """Extract filename from URL with error handling."""
-    try:
-        response = requests.get(url, stream=True)
-        response.raise_for_status()
-        if 'content-disposition' in response.headers:
-            return re.findall('filename="?([^"]+)"?', response.headers['content-disposition'])[0]
-        return os.path.basename(urlparse(url).path)
-    except Exception as e:
-        print(f"Error getting filename: {e}")
-        return "downloaded_model"
-
-def get_supported_extensions():
-    """Return supported model extensions."""
-    return (".ckpt", ".safetensors", ".pt", ".pth")
-
-# ---------------------- MODEL CONVERSION CORE ----------------------
-class ConversionHistory:
-    """Track conversion attempts and provide optimization suggestions."""
-    def __init__(self, history_file="conversion_history.json"):
-        self.history_file = history_file
-        self.history = self._load_history()
-    
-    def _load_history(self):
-        try:
-            with open(self.history_file, 'r') as f:
-                return json.load(f)
-        except:
-            return []
-            
-    def add_entry(self, model_path, settings, success, message):
-        entry = {
-            "timestamp": datetime.now().isoformat(),
-            "model": model_path,
-            "settings": settings,
-            "success": success,
-            "message": message
-        }
-        self.history.append(entry)
-        self._save_history()
-    
-    def get_optimization_suggestions(self, model_path):
-        """Generate suggestions based on conversion history."""
-        suggestions = []
-        for entry in self.history:
-            if entry["model"] == model_path and not entry["success"]:
-                suggestions.append(f"Previous failure: {entry['message']}")
-        return suggestions
-
-def convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, use_xformers, hf_token, orgs_name, model_name, make_private, output_widget):
-    """Main conversion logic with error handling."""
-    history = ConversionHistory()
+def determine_load_checkpoint(model_to_load):
+    """Determines if the model to load is a checkpoint or a Diffusers model."""
+    if model_to_load.endswith('.ckpt') or model_to_load.endswith('.safetensors'):
+        return True
+    elif os.path.isdir(model_to_load):
+        required_folders = {"unet", "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "scheduler", "vae"}
+        if required_folders.issubset(set(os.listdir(model_to_load))) and os.path.isfile(os.path.join(model_to_load, "model_index.json")):
+            return False
+    return None
+
+def increment_filename(filename):
+    """Increments the filename to avoid overwriting existing files."""
+    base, ext = os.path.splitext(filename)
+    counter = 1
+    while os.path.exists(filename):
+        filename = f"{base}({counter}){ext}"
+        counter += 1
+    return filename
+
+def create_model_repo(api, user, orgs_name, model_name, make_private=False):
+    """Creates a Hugging Face model repository if it doesn't exist."""
+    repo_id = f"{orgs_name}/{model_name.strip()}" if orgs_name else f"{user['name']}/{model_name.strip()}"
     try:
-        # Conversion steps here
-        return "Conversion successful!"
-    except Exception as e:
-        history.add_entry(model_to_load, locals(), False, str(e))
-        return f"❌ Error: {str(e)}"
+        validate_repo_id(repo_id)
+        api.create_repo(repo_id=repo_id, repo_type="model", private=make_private)
+        print(f"Model repo '{repo_id}' created.")
+    except HfHubHTTPError:
+        print(f"Model repo '{repo_id}' already exists.")
+
+    return repo_id
+
+# ---------------------- MODEL LOADING AND CONVERSION ----------------------
+def load_sdxl_model(model_to_load, is_load_checkpoint, load_dtype):
+    """Loads the SDXL model from a checkpoint or Diffusers model."""
+    model_load_message = "checkpoint" if is_load_checkpoint else "Diffusers" + (" as fp16" if load_dtype == torch.float16 else "")
+    print(f"Loading {model_load_message}: {model_to_load}")
+
+    if is_load_checkpoint:
+        return load_from_sdxl_checkpoint(model_to_load)
+    else:
+        return load_sdxl_from_diffusers(model_to_load, load_dtype)
+
+def load_from_sdxl_checkpoint(model_to_load):
+    """Loads the SDXL model components from a checkpoint file."""
+    # Implement loading logic here
+    text_encoder1, text_encoder2, vae, unet = None, None, None, None
+    # Example loading logic (replace with actual loading code)
+    # text_encoder1, text_encoder2, vae, unet = sdxl_model_util.load_models_from_sdxl_checkpoint("sdxl_base_v1-0", model_to_load, "cpu")
+    print(f"Loaded from checkpoint: {model_to_load}")
+    return text_encoder1, text_encoder2, vae, unet
+
+def load_sdxl_from_diffusers(model_to_load, load_dtype):
+    """Loads an SDXL model from a Diffusers model directory."""
+    pipeline = StableDiffusionXLPipeline.from_pretrained(model_to_load, torch_dtype=load_dtype)
+    text_encoder1 = pipeline.text_encoder
+    text_encoder2 = pipeline.text_encoder_2
+    vae = pipeline.vae
+    unet = pipeline.unet
+
+    return text_encoder1, text_encoder2, vae, unet
+
+def convert_and_save_sdxl_model(model_to_load, is_save_checkpoint, loaded_model_data, save_dtype):
+    """Converts and saves the SDXL model as either a checkpoint or a Diffusers model."""
+    text_encoder1, text_encoder2, vae, unet = loaded_model_data
+    if is_save_checkpoint:
+        save_sdxl_as_checkpoint(model_to_load, text_encoder1, text_encoder2, vae, unet, save_dtype)
+    else:
+        save_sdxl_as_diffusers(model_to_load, text_encoder1, text_encoder2, vae, unet, save_dtype)
+
+def save_sdxl_as_checkpoint(model_to_save, text_encoder1, text_encoder2, vae, unet, save_dtype):
+    """Saves the SDXL model components as a checkpoint file."""
+    # Implement saving logic here
+    print(f"Model saved as checkpoint: {model_to_save}")
+
+def save_sdxl_as_diffusers(model_to_save, text_encoder1, text_encoder2, vae, unet, save_dtype):
+    """Saves the SDXL model as a Diffusers model."""
+    pipeline = StableDiffusionXLPipeline(
+        vae=vae,
+        text_encoder=text_encoder1,
+        text_encoder_2=text_encoder2,
+        unet=unet
+    )
+    pipeline.save_pretrained(model_to_save)
+    print(f"Model saved as Diffusers format: {model_to_save}")
+
+# ---------------------- UPLOAD FUNCTION ----------------------
+def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_private):
+    """Uploads a model to the Hugging Face Hub."""
+    login(hf_token, add_to_git_credential=True)
+    api = HfApi()
+    user = api.whoami(hf_token)
+    model_repo = create_model_repo(api, user, orgs_name, model_name, make_private)
+
+    # Upload logic here
+    print(f"Model uploaded to: https://huggingface.co/{model_repo}")
 
 # ---------------------- GRADIO INTERFACE ----------------------
-# ---------------------- GRADIO INTERFACE ----------------------
-# Temporarily disabled theme configuration
-# def build_theme(theme_name, font):
-#     """Create accessible theme with dynamic settings."""
-#     return gr.themes.Default().set(
-#         primary_hue="violet" if "dark" in theme_name else "indigo",
-#         font=[font, "ui-sans-serif", "sans-serif"],
-#         button_primary_background="*primary_300",
-#         button_primary_text_color="white",
-#         background_fill="*neutral_50" if "light" in theme_name else "*neutral_950"
-#     )
-
-with gr.Blocks(
-    css="""
-    .single-column {max-width: 800px; margin: 0 auto;}
-    .output-panel {background: rgba(0,0,0,0.05); padding: 20px; border-radius: 8px;}
-    """,
-    # theme=build_theme("dark", "Arial")  # Theme disabled temporarily
-) as demo:
-    
-    # Accessibility Controls
-    with gr.Accordion("♿ Accessibility Settings", open=False):
-        with gr.Row():
-            theme_selector = gr.Dropdown(
-                ["Dark Mode", "Light Mode", "High Contrast"],
-                label="Color Theme",
-                value="Dark Mode"
-            )
-            font_selector = gr.Dropdown(
-                ["Arial", "OpenDyslexic", "Comic Neue"], 
-                label="Font Choice",
-                value="Arial"
-            )
-            font_size = gr.Slider(12, 24, value=16, label="Font Size (px)")
-
-    # Main Content
-    with gr.Column(elem_classes="single-column"):
-        gr.Markdown("""
-        # 🎨 SDXL Model Converter
-        Convert models between formats with accessibility in mind!
-        
-        ### Features:
-        - 🧠 Memory-efficient conversions
-        - ♿ Dyslexia-friendly fonts
-        - 🌓 Dark/Light modes
-        - 🤗 HF Hub integration
-        """)
-        
-        # Input Fields
-        model_to_load = gr.Textbox(label="Model Path/URL")
-        save_precision_as = gr.Dropdown(["float32", "float16"], label="Precision")
-        
-        with gr.Row():
-            epoch = gr.Number(label="Epoch", value=0)
-            global_step = gr.Number(label="Global Step", value=0)
-
-        # Conversion Button
-        convert_btn = gr.Button("Convert", variant="primary")
-        
-        # Output Panel
-        output = gr.Markdown(elem_classes="output-panel")
-
-# ---------------------- MAIN EXECUTION ----------------------
-if __name__ == "__main__":
-    demo.launch(share=True)
+def main(model_to_load, save_precision_as, epoch, global_step, reference_model, output_path, fp16, hf_token, orgs_name, model_name, make_private):
+    """Main function orchestrating the entire process."""
+    load_dtype = get_save_dtype(save_precision_as)
+    is_load_checkpoint = determine_load_checkpoint(model_to_load)
+    is_save_checkpoint = not is_load_checkpoint
+
+    loaded_model_data = load_sdxl_model(model_to_load, is_load_checkpoint, load_dtype)
+    convert_and_save_sdxl_model(model_to_load, is_save_checkpoint, loaded_model_data, load_dtype)
+    upload_to_huggingface(output_path, hf_token, orgs_name, model_name, make_private)
+
+    return "Conversion and upload completed successfully!"
+
+with gr.Blocks() as demo:
+    model_to_load = gr.Textbox(label="Model to Load (Checkpoint or Diffusers)", placeholder="Path to model")
+    save_precision_as = gr.Dropdown(choices=["fp16", "bf16", "float"], label="Save Precision As")
+    epoch = gr.Number(value=0, label="Epoch to Write (Checkpoint)")
+    global_step = gr.Number(value=0, label="Global Step to Write (Checkpoint)")
+    reference_model = gr.Textbox(label="Reference Diffusers Model", placeholder="e.g., stabilityai/stable-diffusion-xl-base-1.0")
+    output_path = gr.Textbox(label="Output Path", value="/content/output")
+    hf_token = gr.Textbox(label="Hugging Face Token", placeholder="Your Hugging Face write token")
+    orgs_name = gr.Textbox(label="Organization Name (Optional)", placeholder="Your organization name")
+    model_name = gr.Textbox(label="Model Name", placeholder="The name of your model on Hugging Face")
+    make_private = gr.Checkbox(label="Make Repository Private", value=False)
+
+    convert_button = gr.Button("Convert and Upload")
+    output = gr.Markdown()
+
+    convert_button.click(fn=main, inputs=[model_to_load, save_precision_as, epoch, global_step, reference_model, output_path, fp16, hf_token, orgs_name, model_name, make_private], outputs=output)
+
+demo.launch()