Update app.py

app.py

@@ -1,17 +1,15 @@
 import gradio as gr
 import numpy as np
-import torch, random, json, spaces, time
+import torch, random, json, spaces
 from ulid import ULID
 from diffsynth.pipelines.qwen_image import (
     QwenImagePipeline, ModelConfig,
     QwenImageUnit_Image2LoRAEncode, QwenImageUnit_Image2LoRADecode
 )
 from safetensors.torch import save_file
-import torch
 from PIL import Image
-from utils import repo_utils, image_utils, prompt_utils
-
 
+# from utils import repo_utils, image_utils, prompt_utils
 # repo_utils.clone_repo_if_not_exists("git clone https://huggingface.co/DiffSynth-Studio/General-Image-Encoders", "app/repos")
 # repo_utils.clone_repo_if_not_exists("https://huggingface.co/apple/starflow", "app/models")
 
@@ -19,6 +17,7 @@ URL_PUBLIC = "https://huggingface.co/spaces/AiSudo/Qwen-Image-to-LoRA/blob/main"
 DTYPE = torch.bfloat16
 MAX_SEED = np.iinfo(np.int32).max
 
+
 vram_config_disk_offload = {
     "offload_dtype": "disk",
     "offload_device": "disk",
@@ -30,28 +29,28 @@ vram_config_disk_offload = {
     "computation_device": "cuda",
 }
 
-# Load models
+# Load models (LoRA encoder/decoder)
 pipe_lora = QwenImagePipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
     device="cuda",
     model_configs=[
         ModelConfig(
-            download_source="huggingface", 
-            model_id="DiffSynth-Studio/General-Image-Encoders", 
-            origin_file_pattern="SigLIP2-G384/model.safetensors", 
-            **vram_config_disk_offload
+            download_source="huggingface",
+            model_id="DiffSynth-Studio/General-Image-Encoders",
+            origin_file_pattern="SigLIP2-G384/model.safetensors",
+            **vram_config_disk_offload,
         ),
         ModelConfig(
-            download_source="huggingface", 
-            model_id="DiffSynth-Studio/General-Image-Encoders", 
-            origin_file_pattern="DINOv3-7B/model.safetensors", 
-            **vram_config_disk_offload
+            download_source="huggingface",
+            model_id="DiffSynth-Studio/General-Image-Encoders",
+            origin_file_pattern="DINOv3-7B/model.safetensors",
+            **vram_config_disk_offload,
         ),
         ModelConfig(
-            download_source="huggingface", 
-            model_id="DiffSynth-Studio/Qwen-Image-i2L", 
-            origin_file_pattern="Qwen-Image-i2L-Style.safetensors", 
-            **vram_config_disk_offload
+            download_source="huggingface",
+            model_id="DiffSynth-Studio/Qwen-Image-i2L",
+            origin_file_pattern="Qwen-Image-i2L-Style.safetensors",
+            **vram_config_disk_offload,
         ),
     ],
     processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Edit", origin_file_pattern="processor/"),
@@ -69,6 +68,7 @@ vram_config = {
     "computation_device": "cuda",
 }
 
+# Load image generation pipeline
 pipe_imagen = QwenImagePipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
     device="cuda",
@@ -82,32 +82,101 @@ pipe_imagen = QwenImagePipeline.from_pretrained(
 )
 
 
+def read_file(path: str) -> str:
+    with open(path, "r", encoding="utf-8") as f:
+        return f.read()
+
+
+def show_user(profile: gr.OAuthProfile | None):
+    """
+    Displays who is logged in (or nothing if not logged in).
+    Works in Spaces OAuth; locally uses hf auth login if available.
+    """
+    if profile is None:
+        return ""
+    username = getattr(profile, "username", None) or "unknown"
+    return f"✅ Signed in as **{username}**"
+
+
 @spaces.GPU
 def generate_lora(
     input_images,
+    profile: gr.OAuthProfile | None = None,
+    oauth_token: gr.OAuthToken | None = None,
     progress=gr.Progress(track_tqdm=True),
 ):
+    """
+    - Always generates and saves LoRA locally under ./loras/
+    - If user is signed in (OAuth), also uploads to the user's *own* Hub repo.
+    """
+    import os
+    from huggingface_hub import HfApi
 
     ulid = str(ULID()).lower()[:12]
     print(f"ulid: {ulid}")
 
     if not input_images:
-        print("images are empty.")
-        return False
+        return (
+            "",
+            gr.update(value="⚠️ Please upload at least 1 image."),
+            gr.update(interactive=False),
+            gr.update(interactive=False, link=""),
+        )
 
+    # Gradio Gallery returns list of (filepath, metadata)
     input_images = [Image.open(filepath).convert("RGB") for filepath, _ in input_images]
-    
+
     # Model inference
     with torch.no_grad():
         embs = QwenImageUnit_Image2LoRAEncode().process(pipe_lora, image2lora_images=input_images)
         lora = QwenImageUnit_Image2LoRADecode().process(pipe_lora, **embs)["lora"]
 
     lora_name = f"{ulid}.safetensors"
+    os.makedirs("loras", exist_ok=True)
     lora_path = f"loras/{lora_name}"
-
     save_file(lora, lora_path)
 
-    return lora_name, gr.update(interactive=True, value=lora_path), gr.update(interactive=True)
+    # Default: local-only message (still lets user generate images from local LoRA)
+    hub_url = ""
+    hub_markdown = "✅ LoRA generated locally. Sign in to upload it to your Hugging Face account."
+
+    # Upload to the signed-in user's own account if available
+    if profile is not None and oauth_token is not None and getattr(oauth_token, "token", None):
+        try:
+            username = getattr(profile, "username", None) or ""
+            if not username:
+                raise ValueError("Could not read username from OAuth profile.")
+
+            api = HfApi(token=oauth_token.token)
+
+            # Create / reuse a user repo (model repo recommended for LoRAs)
+            # Change name if you want:
+            repo_id = f"{username}/qwen-image-loras"
+            api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True)
+
+            api.upload_file(
+                path_or_fileobj=lora_path,
+                path_in_repo=f"loras/{lora_name}",
+                repo_id=repo_id,
+                repo_type="model",
+                commit_message=f"Add LoRA: {lora_name}",
+            )
+
+            hub_url = f"https://huggingface.co/{repo_id}/blob/main/loras/{lora_name}"
+            hub_markdown = f"✅ **Uploaded to your account:** {hub_url}"
+
+        except Exception as e:
+            print(f"Error uploading to user repo: {e}")
+            hub_markdown = f"⚠️ Upload failed (still saved locally): `{str(e)}`"
+
+    # Enable image generation button regardless
+    return (
+        lora_name,
+        gr.update(value=hub_markdown),
+        gr.update(interactive=True),
+        gr.update(interactive=bool(hub_url), link=hub_url),
+    )
+
 
 @spaces.GPU
 def generate_image(
@@ -122,14 +191,18 @@ def generate_image(
     num_inference_steps=8,
     progress=gr.Progress(track_tqdm=True),
 ):
+    if not lora_name:
+        return None, seed
+
     lora_path = f"loras/{lora_name}"
+
     pipe_imagen.clear_lora()
     pipe_imagen.load_lora(pipe_imagen.dit, lora_path)
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator().manual_seed(seed)
+    # generator = torch.Generator().manual_seed(seed)
 
     output_image = pipe_imagen(
         prompt=prompt,
@@ -144,13 +217,6 @@ def generate_image(
 
     return output_image, seed
 
-    return True
-
-
-def read_file(path: str) -> str:
-    with open(path, 'r', encoding='utf-8') as f:
-        content = f.read()
-    return content
 
 # Enhanced Apple-style CSS - more minimalist and clean
 css = """
@@ -302,20 +368,6 @@ css = """
     background: #30a14a !important;
 }
 
-/* Checkbox styling */
-.gradio-container .gr-checkbox {
-    background: #ffffff !important;
-    border: 2px solid #d2d2d7 !important;
-    border-radius: 6px !important;
-    width: 24px !important;
-    height: 24px !important;
-}
-
-.gradio-container .gr-checkbox.checked {
-    background: #007aff !important;
-    border-color: #007aff !important;
-}
-
 /* Examples section */
 .gradio-container .gr-examples {
     background: #f2f2f7 !important;
@@ -324,143 +376,94 @@ css = """
     border: none !important;
 }
 
-/* Label styling */
-.gradio-container .gr-label {
-    font-weight: 600 !important;
-    color: #1d1d1f !important;
-    font-size: 17px !important;
-    margin-bottom: 8px !important;
-    letter-spacing: -.022em !important;
-}
-
 /* Mobile responsiveness */
 @media (max-width: 768px) {
-    #col-container {
-        padding: 20px 16px !important;
-        max-width: 100% !important;
-    }
-    
-    .gradio-container h1 {
-        font-size: 2rem !important;
-        margin-bottom: 12px !important;
-    }
-    
-    .gradio-container .subtitle {
-        font-size: 1.1rem !important;
-    }
-    
-    .section-card {
-        padding: 24px !important;
-        margin-bottom: 24px !important;
-    }
-    
-    .gradio-container .gr-button {
-        padding: 14px 28px !important;
-        font-size: 16px !important;
-    }
-    
-    .gradio-container .gr-gallery {
-        height: 200px !important;
-        columns: 2 !important;
-    }
-    
-    .gradio-container .gr-row {
-        flex-direction: column !important;
-        gap: 20px !important;
-    }
+    #col-container { padding: 20px 16px !important; max-width: 100% !important; }
+    .gradio-container h1 { font-size: 2rem !important; margin-bottom: 12px !important; }
+    .gradio-container .subtitle { font-size: 1.1rem !important; }
+    .section-card { padding: 24px !important; margin-bottom: 24px !important; }
+    .gradio-container .gr-button { padding: 14px 28px !important; font-size: 16px !important; }
+    .gradio-container .gr-gallery { height: 200px !important; columns: 2 !important; }
+    .gradio-container .gr-row { flex-direction: column !important; gap: 20px !important; }
 }
 
 @media (max-width: 480px) {
-    .gradio-container h1 {
-        font-size: 1.75rem !important;
-    }
-    
-    .section-card {
-        padding: 20px !important;
-    }
-    
-    .gradio-container .gr-gallery {
-        height: 180px !important;
-        columns: 1 !important;
-    }
+    .gradio-container h1 { font-size: 1.75rem !important; }
+    .section-card { padding: 20px !important; }
+    .gradio-container .gr-gallery { height: 180px !important; columns: 1 !important; }
 }
 
-/* Remove all animations for Apple-like snappiness */
-.gradio-container * {
-    transition: background 0.15s ease, border-color 0.15s ease, color 0.15s ease !important;
-}
-
-/* Hide unnecessary elements */
+/* Hide gradio header/footer */
 .gradio-container .gr-footer,
 .gradio-container .gr-header {
     display: none !important;
 }
-
-/* Clean scrollbar */
-.gradio-container ::-webkit-scrollbar {
-    width: 3px !important;
-}
-
-.gradio-container ::-webkit-scrollbar-track {
-    background: transparent !important;
-}
-
-.gradio-container ::-webkit-scrollbar-thumb {
-    background: #d2d2d7 !important;
-    border-radius: 3px !important;
-}
-
-.gradio-container ::-webkit-scrollbar-thumb:hover {
-    background: #007aff !important;
-}
 """
 
 
-
-with open('examples/0_examples.json', 'r') as file: examples = json.load(file)
+# Load examples
+with open("examples/0_examples.json", "r") as file:
+    examples = json.load(file)
 print(examples)
 
-# Gradio 6 syntax - no parameters in Blocks()
+
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
-        with gr.Column():
-            # Ultra-minimalist Apple-style header
-            gr.HTML("""
-            <div style="text-align: center; max-width: 600px; margin: 0 auto;">
+        # Header
+        gr.HTML(
+            """
+            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
                 <h1>Qwen Image to LoRA</h1>
                 <p class="subtitle">Generate custom LoRA models from your images</p>
                 <p style="font-size: 14px; color: #86868b; margin-top: 16px;">
-                    Demo by <a href="https://aisudo.com/" target="_blank" style="color: #007aff; text-decoration: none;">AiSudo</a> • 
+                    Demo by <a href="https://aisudo.com/" target="_blank" style="color: #007aff; text-decoration: none;">AiSudo</a> •
                     <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #007aff; text-decoration: none;">Built with anycoder</a>
                 </p>
             </div>
-            """)
-        
+            """
+        )
+
+        # ✅ Hugging Face Login Button
+        with gr.Row():
+            with gr.Column(scale=1):
+                login_btn = gr.LoginButton(
+                    value="Sign in with Hugging Face",
+                    logout_value="Logout ({})",
+                    variant="huggingface",
+                    size="lg",
+                )
+            with gr.Column(scale=3):
+                whoami = gr.Markdown(value="", elem_id="whoami")
+
         with gr.Row():
             with gr.Column(elem_classes=["section-card"]):
                 input_images = gr.Gallery(
-                    label="Input Images", 
+                    label="Input Images",
                     file_types=["image"],
-                    show_label=True, 
-                    columns=2, 
-                    object_fit="cover", 
-                    height=250)
-                
+                    show_label=True,
+                    columns=2,
+                    object_fit="cover",
+                    height=250,
+                )
                 lora_button = gr.Button("Generate LoRA", size="lg")
 
             with gr.Column(elem_classes=["section-card"]):
                 lora_name = gr.Textbox(
-                    label="Generated LoRA", 
-                    lines=2, 
+                    label="Generated LoRA",
+                    lines=2,
                     interactive=False,
-                    placeholder="Your LoRA will appear here..."
+                    placeholder="Your LoRA will appear here...",
                 )
-                lora_download = gr.DownloadButton(
-                    label="Download LoRA", 
+                hub_link = gr.Markdown(value="", label="Hub Link")
+
+                # This becomes clickable only after upload
+                lora_download = gr.Button(
+                    value="View on Hub",
                     interactive=False,
-                    size="lg"
+                    size="lg",
+                    link="",
                 )
-        
+
         with gr.Column(elem_classes=["section-card"]) as imagen_container:
             gr.Markdown("### Generate Images")
             with gr.Row():
@@ -471,15 +474,15 @@ with gr.Blocks() as demo:
                         placeholder="Describe what you want to generate...",
                         value="a person in a fishing boat.",
                     )
-                    
-                    imagen_button = gr.Button("Generate Image", interactive=False, size="lg")    
-                    
+
+                    imagen_button = gr.Button("Generate Image", interactive=False, size="lg")
+
                     with gr.Accordion("Settings", open=False):
                         negative_prompt = gr.Textbox(
                             label="Negative Prompt",
                             lines=1,
                             placeholder="What to avoid...",
-                            value="blurry, low quality"
+                            value="blurry, low quality",
                         )
                         num_inference_steps = gr.Slider(
                             label="Steps",
@@ -501,7 +504,7 @@ with gr.Blocks() as demo:
                                 minimum=512,
                                 maximum=1280,
                                 step=32,
-                                value=768, 
+                                value=768,
                             )
                             height = gr.Slider(
                                 label="Height",
@@ -520,26 +523,23 @@ with gr.Blocks() as demo:
                         randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
 
                 with gr.Column():
-                    output_image = gr.Image(
-                        label="Generated Image", 
-                        height=350
-                    )
-                    
-        gr.Examples(
-            examples=examples, 
-            inputs=[input_images],
-            label="Examples"
-        )
+                    output_image = gr.Image(label="Generated Image", height=350)
+
+        gr.Examples(examples=examples, inputs=[input_images], label="Examples")
         gr.Markdown(read_file("static/footer.md"))
 
+    # Login click shows the user
+    login_btn.click(fn=show_user, inputs=[login_btn], outputs=[whoami], api_visibility="public")
+
+    # Generate LoRA (auto-uploads to user account if signed in)
     lora_button.click(
         fn=generate_lora,
-        inputs=[
-            input_images
-        ],
-        outputs=[lora_name, lora_download, imagen_button],
-        api_visibility="public"
+        inputs=[input_images],
+        outputs=[lora_name, hub_link, imagen_button, lora_download],
+        api_visibility="public",
     )
+
+    # Generate Image
     imagen_button.click(
         fn=generate_image,
         inputs=[
@@ -554,33 +554,20 @@ with gr.Blocks() as demo:
             num_inference_steps,
         ],
         outputs=[output_image, seed],
-        api_visibility="public"
+        api_visibility="public",
     )
 
 
 if __name__ == "__main__":
-    # Gradio 6 syntax - all launch parameters go here
     demo.launch(
-        css=css,  # Moved from Blocks() to launch() - Gradio 6 syntax
-        mcp_server=True, 
+        css=css,  # Gradio 6: pass css here
+        mcp_server=True,
         theme=gr.themes.Base(
             primary_hue="blue",
             secondary_hue="gray",
             neutral_hue="gray",
-            font=[
-                gr.themes.GoogleFont("Inter"),
-                "ui-sans-serif",
-                "system-ui",
-                "sans-serif"
-            ],
-            font_mono=[
-                gr.themes.GoogleFont("JetBrains Mono"),
-                "ui-monospace",
-                "Consolas",
-                "monospace"
-            ]
+            font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
+            font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "Consolas", "monospace"],
         ),
-        footer_links=[
-            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
-        ]
-    )
+        footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}],
+    )