Update app.py

app.py

@@ -3,23 +3,68 @@ from transformers import CLIPTextModel, CLIPTokenizer
 import torch
 import gradio as gr
 import spaces
+from huggingface_hub import hf_hub_download
+import os
+import requests
+import hashlib
+from pathlib import Path
+import re
 
+# Default LoRA for fallback
+DEFAULT_LORA = "OedoSoldier/detail-tweaker-lora"
+LORA_CACHE_DIR = "lora_cache"
 
-lora1 = "OedoSoldier/detail-tweaker-lora"
+def download_lora(url):
+    """Download LoRA file from Civitai URL and cache it locally"""
+    # Create cache directory if it doesn't exist
+    os.makedirs(LORA_CACHE_DIR, exist_ok=True)
+    
+    # Generate a filename from the URL
+    url_hash = hashlib.md5(url.encode()).hexdigest()
+    local_path = os.path.join(LORA_CACHE_DIR, f"{url_hash}.safetensors")
+    
+    # If file already exists in cache, return the path
+    if os.path.exists(local_path):
+        return local_path
+    
+    # Download the file
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        
+        # Get the total file size
+        total_size = int(response.headers.get('content-length', 0))
+        
+        # Download and save the file
+        with open(local_path, 'wb') as f:
+            if total_size == 0:
+                f.write(response.content)
+            else:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+        
+        return local_path
+    except Exception as e:
+        print(f"Error downloading LoRA: {str(e)}")
+        return None
+
+def is_civitai_url(url):
+    """Check if the URL is a valid Civitai download URL"""
+    return bool(re.match(r'https?://civitai\.com/api/download/models/\d+', url))
 
 @spaces.GPU
-def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_scale=7.0,model="Real6.0",num_images=1, width=512, height=512):
-        
+def generate_image(prompt, negative_prompt, lora_url, num_inference_steps=30, guidance_scale=7.0, 
+                  model="Real6.0", num_images=1, width=512, height=512):
+    
     if model == "Real5.0":
         model_id = "SG161222/Realistic_Vision_V5.0_noVAE"
-        
     elif model == "Real5.1":
         model_id = "SG161222/Realistic_Vision_V5.1_noVAE"
-        
     else:
         model_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
 
-
+    # Initialize models
     vae = AutoencoderKL.from_pretrained(
         model_id,
         subfolder="vae"
@@ -47,13 +92,30 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
         vae=vae
     ).to("cuda")
 
-    pipe.load_lora_weights(lora1)
+    # Load LoRA weights
+    try:
+        if lora_url and lora_url.strip():
+            if is_civitai_url(lora_url):
+                # Download and load Civitai LoRA
+                lora_path = download_lora(lora_url)
+                if lora_path:
+                    pipe.load_lora_weights(lora_path)
+                else:
+                    pipe.load_lora_weights(DEFAULT_LORA)
+            # If it's a HuggingFace repo path
+            elif '/' in lora_url and not lora_url.startswith('http'):
+                pipe.load_lora_weights(lora_url)
+            else:
+                pipe.load_lora_weights(DEFAULT_LORA)
+        else:
+            pipe.load_lora_weights(DEFAULT_LORA)
+    except Exception as e:
+        print(f"Error loading LoRA weights: {str(e)}")
+        pipe.load_lora_weights(DEFAULT_LORA)
     
     if model == "Real6.0":
         pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
 
-    
-
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(
         pipe.scheduler.config,
         algorithm_type="dpmsolver++",
@@ -79,7 +141,6 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
     prompt_embeds = text_encoder(text_inputs.input_ids)[0]
     negative_prompt_embeds = text_encoder(negative_text_inputs.input_ids)[0]
 
-    
     # Generate the image
     result = pipe(
         prompt_embeds=prompt_embeds,
@@ -94,6 +155,17 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
     
     return result.images
 
+def clean_lora_cache():
+    """Clean the LoRA cache directory"""
+    if os.path.exists(LORA_CACHE_DIR):
+        for file in os.listdir(LORA_CACHE_DIR):
+            file_path = os.path.join(LORA_CACHE_DIR, file)
+            try:
+                if os.path.isfile(file_path):
+                    os.unlink(file_path)
+            except Exception as e:
+                print(f"Error deleting {file_path}: {str(e)}")
+
 title = """<h1 align="center">ProFaker</h1>"""
 # Create the Gradio interface
 with gr.Blocks() as demo:
@@ -112,28 +184,34 @@ with gr.Blocks() as demo:
                 info="Enter what you don't want in Image...",
                 lines=3
             )
+            lora_input = gr.Textbox(
+                label="LoRA URL/Path",
+                info="Enter Civitai download URL or HuggingFace path (e.g., 'username/model-name')",
+                value=DEFAULT_LORA
+            )
+            clear_cache = gr.Button("Clear LoRA Cache")
             generate_button = gr.Button("Generate Image")
+            
             with gr.Accordion("Advanced Options", open=False):
-                
                 model = gr.Dropdown(
                     choices=["Real6.0","Real5.1","Real5.0"],
                     value="Real6.0",
                     label="Model",
                 )
                 
-                num_images = gr.Slider(  # New slider for number of images
-                        minimum=1,
-                        maximum=4,
-                        value=1,
-                        step=1,
-                        label="Number of Images to Generate"
+                num_images = gr.Slider(
+                    minimum=1,
+                    maximum=4,
+                    value=1,
+                    step=1,
+                    label="Number of Images to Generate"
                 )
                 width = gr.Slider(
-                        minimum=256,
-                        maximum=1024,
-                        value=512,
-                        step=64,
-                        label="Image Width"
+                    minimum=256,
+                    maximum=1024,
+                    value=512,
+                    step=64,
+                    label="Image Width"
                 )
                 height = gr.Slider(
                     minimum=256,
@@ -156,6 +234,7 @@ with gr.Blocks() as demo:
                     step=0.5,
                     label="Guidance Scale"
                 )
+        
         with gr.Column():
             # Output component
             gallery = gr.Gallery(
@@ -165,13 +244,16 @@ with gr.Blocks() as demo:
                 columns=2,
                 rows=2
             )
-            
     
     # Connect the interface to the generation function
     generate_button.click(
         fn=generate_image,
-        inputs=[prompt, negative_prompt, steps_slider, guidance_slider, model, num_images, width, height],
+        inputs=[prompt, negative_prompt, lora_input, steps_slider, guidance_slider, 
+                model, num_images, width, height],
         outputs=gallery
     )
+    
+    # Connect clear cache button
+    clear_cache.click(fn=clean_lora_cache)
 
-demo.queue(max_size=10).launch(share=False)
+demo.queue(max_size=10).launch(share=False)