Update app.py

app.py

@@ -1,63 +1,89 @@
 import torch
 from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
 from tqdm.auto import tqdm
-from huggingface_hub import cached_download, hf_hub_url
+from huggingface_hub import hf_hub_url, login, HfApi, create_repo
 import os
+import traceback
+from peft import PeftModel
+import gradio as gr
 
 def display_image(image):
-    """
-    Replace this with your actual image display logic.
-    """
-    image.show()
+    """Display the generated image."""
+    return image 
 
-def load_and_merge_lora(base_model_id, lora_id, lora_weight_name, lora_adapter_name):
-    try: 
+def load_and_merge_lora(base_model_id, lora_id, lora_adapter_name):
+    try:
         pipe = DiffusionPipeline.from_pretrained(
-            base_model_id, 
-            torch_dtype=torch.float16, 
-            scheduler=DPMSolverMultistepScheduler.from_config(
-                pipe.scheduler.config),
+            base_model_id,
+            torch_dtype=torch.float16,
             variant="fp16",
             use_safetensors=True,
         ).to("cuda")
 
-        lora_url = hf_hub_url(lora_id, revision="main", filename=lora_weight_name)
-        lora_path = cached_download(lora_url) 
+        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+            pipe.scheduler.config
+        )
+
+        # Get the UNet model from the pipeline
+        unet = pipe.unet
 
-        with tqdm(desc="Loading LoRA weights", unit="step") as pbar:
-            pipe.load_lora_weights(
-                lora_path,
-                weight_name=lora_weight_name,
-                adapter_name=lora_adapter_name,
-                progress_callback=lambda step, max_steps: pbar.update(1)
-            )
+        # Apply PEFT to the UNet model 
+        unet = PeftModel.from_pretrained(
+            unet, 
+            lora_id, 
+            torch_dtype=torch.float16, 
+            adapter_name=lora_adapter_name
+        )
+
+        # Replace the original UNet in the pipeline with the PEFT-loaded one
+        pipe.unet = unet
 
         print("LoRA merged successfully!")
         return pipe
+
     except Exception as e:
-        print(f"Error merging LoRA: {e}")
+        error_msg = traceback.format_exc()  
+        print(f"Error merging LoRA: {e}\n\nFull traceback saved to errors.txt")
+
+        with open("errors.txt", "w") as f:
+            f.write(error_msg)
+
         return None
 
-def save_merged_model(pipe, save_path):
-    """Saves the merged model to the specified path."""
+def save_merged_model(pipe, save_path, push_to_hub=False, hf_token=None):
+    """Saves and optionally pushes the merged model to Hugging Face Hub."""
     try:
         pipe.save_pretrained(save_path)
         print(f"Merged model saved successfully to: {save_path}")
-    except Exception as e:
-        print(f"Error saving the merged model: {e}")
 
-if __name__ == "__main__":
-    base_model_id = input("Enter the base model ID: ") 
-    lora_id = input("Enter the LoRA Hugging Face Hub ID: ") 
-    lora_weight_name = input("Enter the LoRA weight file name: ")
-    lora_adapter_name = input("Enter the LoRA adapter name: ")
+        if push_to_hub:
+            if hf_token is None:
+                hf_token = input("Enter your Hugging Face write token: ")
+                login(token=hf_token)
 
-    pipe = load_and_merge_lora(base_model_id, lora_id, lora_weight_name, lora_adapter_name)
+            repo_name = input("Enter the Hugging Face repository name "
+                              "(e.g., your_username/your_model_name): ")
 
-    if pipe:
-        prompt = input("Enter your prompt: ")
-        lora_scale = float(input("Enter the LoRA scale (e.g., 0.9): ")) 
+            # Create the repository if it doesn't exist
+            create_repo(repo_name, token=hf_token, exist_ok=True) 
+
+            api = HfApi()
+            api.upload_folder(
+                folder_path=save_path,
+                repo_id=repo_name,
+                token=hf_token,
+                repo_type="model",
+            )
+            print(f"Model pushed successfully to Hugging Face Hub: {repo_name}")
+
+    except Exception as e:
+        print(f"Error saving/pushing the merged model: {e}")
+
+def generate_and_save(base_model_id, lora_id, lora_adapter_name, prompt, lora_scale, save_path, push_to_hub, hf_token):
+    pipe = load_and_merge_lora(base_model_id, lora_id, lora_adapter_name)
 
+    if pipe:
+        lora_scale = float(lora_scale)
         image = pipe(
             prompt, 
             num_inference_steps=30, 
@@ -65,10 +91,31 @@ if __name__ == "__main__":
             generator=torch.manual_seed(0)
         ).images[0]
 
-        display_image(image)
+        image.save("generated_image.png")
+        print(f"Image saved to: generated_image.png")
 
-        # Ask the user for a directory to save the model
-        save_path = input(
-            "Enter the directory where you want to save the merged model: "
-        )
-        save_merged_model(pipe, save_path) 
+        save_merged_model(pipe, save_path, push_to_hub, hf_token)
+
+        return image, "Image generated and model saved/pushed (if selected)."
+
+iface = gr.Interface(
+    fn=generate_and_save,
+    inputs=[
+        gr.Textbox(label="Base Model ID (e.g., stabilityai/stable-diffusion-xl-base-1.0)"),
+        gr.Textbox(label="LoRA ID (e.g., your_username/your_lora)"),
+        gr.Textbox(label="LoRA Adapter Name"),
+        gr.Textbox(label="Prompt"),
+        gr.Slider(label="LoRA Scale", minimum=0.0, maximum=1.0, value=0.7, step=0.1),
+        gr.Textbox(label="Save Path"),
+        gr.Checkbox(label="Push to Hugging Face Hub"),
+        gr.Textbox(label="Hugging Face Write Token", type="password")
+    ],
+    outputs=[
+        gr.Image(label="Generated Image"),
+        gr.Textbox(label="Status")
+    ],
+    title="LoRA Merger and Image Generator",
+    description="Merge a LoRA with a base Stable Diffusion model and generate images."
+)
+
+iface.launch()