Update app.py

app.py

@@ -1,8 +1,8 @@
 import os
 import torch
 import gradio as gr
+from diffusers import FluxPipeline, DiffusionPipeline
 import spaces
-from diffusers import FluxPipeline, DiffusionPipeline, AutoPipelineForText2Image
 
 # Helper function to get the Hugging Face token securely
 def get_hf_token():
@@ -26,23 +26,23 @@ models = {
     "FLUX.1-schnell": {
         "pipeline_class": FluxPipeline,
         "model_id": "black-forest-labs/FLUX.1-schnell",
-        "config": {"torch_dtype": torch.bfloat16, "variant": "fp16"},
+        "config": {"torch_dtype": torch.bfloat16},
         "description": "**FLUX.1-schnell** is a fast and efficient model designed for quick image generation. It excels at producing high-quality images rapidly, making it ideal for applications where speed is crucial. However, its rapid generation may slightly compromise on the level of detail compared to slower, more meticulous models.",
     },
     "FLUX.1-dev": {
-        "pipeline_class": AutoPipelineForText2Image,
+        "pipeline_class": DiffusionPipeline,
         "model_id": "black-forest-labs/FLUX.1-dev",
         "lora": {
             "repo": "strangerzonehf/Flux-Enrich-Art-LoRA",
             "trigger_word": "enrich art",
         },
-        "config": {"torch_dtype": torch.bfloat16, "variant": "fp16"},
+        "config": {"torch_dtype": torch.bfloat16},
         "description": "**FLUX.1-dev** is a development model that focuses on delivering highly detailed and artistically rich images.",
     },
     "Flux.1-lite-8B-alpha": {
         "pipeline_class": FluxPipeline,
         "model_id": "Freepik/flux.1-lite-8B-alpha",
-        "config": {"torch_dtype": torch.bfloat16, "variant": "fp16"},
+        "config": {"torch_dtype": torch.bfloat16},
         "description": "**Flux.1-lite-8B-alpha** is a lightweight model optimized for efficiency and ease of use.",
     },
 }
@@ -65,7 +65,7 @@ def clear_gpu_memory():
         print(f"Error clearing GPU memory: {e}")
         return f"Error clearing GPU memory: {e}"
 
-@spaces.GPU(duration=80)
+
 def load_model(model_key):
     """Loads a model, clearing GPU memory first if a different model was loaded."""
     global model_load_status
@@ -82,13 +82,12 @@ def load_model(model_key):
         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", {}), low_cpu_mem_usage=True)
+        
+        pipe = pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))
 
         if "lora" in config:
             lora_config = config["lora"]
-            pipe.load_lora_weights(lora_config["repo"], token=_HF_TOKEN, weight_name="FluxEnrichArtLoRA.safetensors")
-            pipe.fuse_lora()
+            pipe.load_lora_weights(lora_config["repo"], token=_HF_TOKEN)
 
         if torch.cuda.is_available():
             pipe.to("cuda")
@@ -100,18 +99,13 @@ def load_model(model_key):
         model_load_status[model_key] = "Failed" # Update load status on error
         return f"Error loading model '{model_key}': {e}"
 
-@spaces.GPU(duration=60)
-def generate_image(model_key, prompt, seed=-1):
-
-    model = loaded_models[model_key]
-
+@spaces.GPU(duration=120)
+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(model_key, {}):
-        prompt = f"{models[model_key]['lora']['trigger_word']}, {prompt}"
-
     with torch.no_grad():
         image = model(prompt=prompt, generator=generator).images[0]
     return image
@@ -130,7 +124,9 @@ def gradio_generate(selected_model, prompt, seed):
         # If still not loaded after attempt, return an error
         return f"Model not loaded. Load status: {model_load_status.get(selected_model, 'Not attempted')}.", None
 
-    image = generate_image(selected_model, prompt, seed)
+    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)
@@ -157,12 +153,18 @@ with gr.Blocks(
         with gr.Column():
             gr.Markdown("# Text-to-Image Generator")
             model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model")
+            #with gr.Row():
+                #load_button = gr.Button("Load Model")
+                # clear_button = gr.Button("Clear GPU Memory") # Removed clear button
+            #load_status = gr.Textbox(label="Model Load Status", interactive=False) # Removed load button
             prompt_textbox = gr.Textbox(label="Enter Text Prompt")
             seed_slider = gr.Slider(minimum=-1, maximum=1000, step=1, value=-1, label="Random Seed (-1 for random)")
             generate_button = gr.Button("Generate Image")
             output_image = gr.Image(label="Generated Image")
             runtime_info_textbox = gr.Textbox(label="Runtime Information", lines=2, interactive=False)
 
+            #load_button.click(gradio_load_model, inputs=[model_dropdown], outputs=[load_status])   # Removed load button click action
+            # clear_button.click(clear_gpu_memory, outputs=[load_status]) # Removed clear button click action
             generate_button.click(gradio_generate, inputs=[model_dropdown, prompt_textbox, seed_slider], outputs=[output_image, runtime_info_textbox])
 
     with gr.Tab("Model Information"):
@@ -174,4 +176,4 @@ with gr.Blocks(
     **Credits**: Created by Ruslan Magana Vsevolodovna. For more information, visit [https://ruslanmv.com/](https://ruslanmv.com/).""")
 
 if __name__ == "__main__":
-    interface.launch(debug=True, share=True)
+    interface.launch(debug=True)