Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import torch
 import gradio as gr
 import spaces
-from diffusers import FluxPipeline, DiffusionPipeline
+from diffusers import FluxPipeline, DiffusionPipeline, AutoPipelineForText2Image
 
 # Helper function to get the Hugging Face token securely
 def get_hf_token():
@@ -30,13 +30,13 @@ models = {
         "description": "**FLUX.1-schnell** is a fast and efficient model designed for quick image generation.",
     },
     "FLUX.1-dev": {
-        "pipeline_class": DiffusionPipeline,
+        "pipeline_class": AutoPipelineForText2Image, # Changed to AutoPipelineForText2Image
         "model_id": "black-forest-labs/FLUX.1-dev",
         "lora": {
             "repo": "strangerzonehf/Flux-Enrich-Art-LoRA",
             "trigger_word": "enrich art",
         },
-        "config": {"torch_dtype": torch.bfloat16},
+        "config": {"torch_dtype": torch.bfloat16, "variant": "fp16"},
         "description": "**FLUX.1-dev** is a development model for highly detailed and artistic images.",
     },
 }
@@ -52,32 +52,33 @@ def download_all_models():
             model_id = config["model_id"]
             pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))
             if "lora" in config:
-                pipeline_class.download_lora_weights(config["lora"]["repo"], token=_HF_TOKEN)
+                # Download LoRA weights separately if needed, but not essential for pre-download
+                pass
             print(f"Model '{model_key}' downloaded successfully.")
         except Exception as e:
             print(f"Error downloading model '{model_key}': {e}")
     print("Model download process complete.")
 
 # Function to load a model when requested
-@spaces.GPU
+#@spaces.GPU # Removed as it's not a standard decorator
 def load_model(model_key):
     if model_key in loaded_models:
         return f"Model '{model_key}' is already loaded."
-
     try:
         config = models[model_key]
         pipeline_class = config["pipeline_class"]
         model_id = config["model_id"]
 
-        pipe = pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))
-
+        pipe = pipeline_class.from_pretrained(
+            model_id, token=_HF_TOKEN, **config.get("config", {}), low_cpu_mem_usage=True # Add low_cpu_mem_usage
+        )
         if "lora" in config:
             lora_config = config["lora"]
-            pipe.load_lora_weights(lora_config["repo"], token=_HF_TOKEN)
+            pipe.load_lora_weights(lora_config["repo"], token=_HF_TOKEN, weight_name="FluxEnrichArtLoRA.safetensors") # Added specific weight file name
+            pipe.fuse_lora()
 
         if torch.cuda.is_available():
             pipe.to("cuda")
-
         loaded_models[model_key] = pipe
         return f"Model '{model_key}' loaded successfully."
     except Exception as e:
@@ -86,12 +87,15 @@ def load_model(model_key):
         return f"Error loading model '{model_key}': {e}"
 
 # Function to generate an image from text
-@spaces.GPU
+#@spaces.GPU # Removed as it's not a standard decorator
 def generate_image(model, prompt, seed=-1):
     generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
     if seed != -1:
         generator = generator.manual_seed(seed)
-
+    
+    if "lora" in models.get(list(loaded_models.keys())[list(loaded_models.values()).index(model)], {}):
+        prompt = f"{models[list(loaded_models.keys())[list(loaded_models.values()).index(model)]]['lora']['trigger_word']}, {prompt}"
+    
     with torch.no_grad():
         image = model(prompt=prompt, generator=generator).images[0]
     return image
@@ -100,10 +104,8 @@ def generate_image(model, prompt, seed=-1):
 def gradio_generate(selected_model, prompt, seed):
     if selected_model not in loaded_models:
         return "Model not loaded. Please load the model first.", None
-
     model = loaded_models[selected_model]
     image = generate_image(model, prompt, seed)
-
     runtime_info = f"Model: {selected_model}\nSeed: {seed}"
     output_path = "generated_image.png"
     image.save(output_path)
@@ -128,14 +130,12 @@ with gr.Blocks() as interface:
             generate_button = gr.Button("Generate Image")
             output_image = gr.Image(label="Generated Image")
             runtime_info_textbox = gr.Textbox(label="Runtime Information", lines=2)
-
             load_button.click(gradio_load_model, inputs=[model_dropdown], outputs=[load_status])
             generate_button.click(
                 fn=gradio_generate,
                 inputs=[model_dropdown, prompt_textbox, seed_slider],
                 outputs=[output_image, runtime_info_textbox],
             )
-
     with gr.Tab("Model Information"):
         with gr.Column():
             for model_key, model_info in models.items():
@@ -144,4 +144,8 @@ with gr.Blocks() as interface:
 
 if __name__ == "__main__":
     download_all_models()
-    interface.launch(debug=True)
+    try:
+        interface.launch(debug=True, share=True) # Enable sharing for public link
+    except gr.Error as e:
+        print(f"Error launching Gradio interface: {e}")
+        print("Please make sure you have the necessary permissions and try again.")