Upload folder using huggingface_hub

app.py

@@ -1,15 +1,14 @@
 import spaces
 import gradio as gr
 import torch
-from diffusers import ZImagePipeline
+from diffusers import ZImagePipeline, AutoPipelineForText2Image
 import os
 from pathlib import Path
 
-# Load the model directly at startup
+# Load the base model directly at startup
 print("Loading Z-Image Turbo model...")
 print("This may take a few minutes on first run while the model downloads...")
 
-
 # Load the pipeline with optimal settings
 pipe = ZImagePipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
@@ -24,21 +23,27 @@ print(f"Model loaded on {device}")
     
 print("Model loaded successfully!")
 
+# Store the current model state
+current_model = "base"
+lora_loaded = False
+
 @spaces.GPU()
 def generate_image(
     prompt,
+    model_choice,
     progress=gr.Progress(track_tqdm=True)
 ):
     """
-    Generate an image using Z-Image Turbo model.
+    Generate an image using Z-Image Turbo model with optional LoRA.
     
     Args:
         prompt: Text description of the desired image
+        model_choice: Either "Base Model" or "Classic Painting LoRA"
     
     Returns:
         Generated PIL Image
     """
-    global pipe
+    global pipe, current_model, lora_loaded
     
     if pipe is None:
         raise gr.Error("Model failed to load on startup. Please restart the application.")
@@ -49,22 +54,69 @@ def generate_image(
     # Determine device
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
-    # Set random seed for reproducibility
-    generator = torch.Generator(device).manual_seed(42)
-    
-    # Generate the image with optimal settings
-    progress(0.1, desc="Generating image...")
+    # Handle model switching
+    progress(0.05, desc="Loading model...")
     
     try:
-        result = pipe(
-            prompt=prompt,
-            negative_prompt=None,
-            height=1024,
-            width=1024,
-            num_inference_steps=9,
-            guidance_scale=0.0,
-            generator=generator,
-        )
+        if model_choice == "Classic Painting LoRA" and current_model != "lora":
+            # Load LoRA weights
+            if not lora_loaded:
+                print("Loading Classic Painting LoRA...")
+                pipe.load_lora_weights(
+                    "renderartist/Classic-Painting-Z-Image-Turbo-LoRA",
+                    adapter_name="classic_painting",
+                    weight_name="pytorch_lora_weights.safetensors"
+                )
+                lora_loaded = True
+            
+            # Set LoRA adapter
+            pipe.set_adapters(["classic_painting"], adapter_weights=[0.8])
+            current_model = "lora"
+            progress(0.15, desc="LoRA loaded, generating image...")
+            
+            # Generate with LoRA settings
+            generator = torch.Generator(device).manual_seed(42)
+            result = pipe(
+                prompt=prompt,
+                negative_prompt=None,
+                height=1024,
+                width=1024,
+                num_inference_steps=9,
+                guidance_scale=0.0,
+                generator=generator,
+            )
+        
+        elif model_choice == "Base Model" and current_model != "base":
+            # Disable LoRA
+            pipe.disable_lora()
+            current_model = "base"
+            progress(0.15, desc="Generating image...")
+            
+            # Generate with base model settings
+            generator = torch.Generator(device).manual_seed(42)
+            result = pipe(
+                prompt=prompt,
+                negative_prompt=None,
+                height=1024,
+                width=1024,
+                num_inference_steps=9,
+                guidance_scale=0.0,
+                generator=generator,
+            )
+        
+        else:
+            # Model already loaded, just generate
+            progress(0.15, desc="Generating image...")
+            generator = torch.Generator(device).manual_seed(42)
+            result = pipe(
+                prompt=prompt,
+                negative_prompt=None,
+                height=1024,
+                width=1024,
+                num_inference_steps=9,
+                guidance_scale=0.0,
+                generator=generator,
+            )
         
         image = result.images[0]
         progress(1.0, desc="Complete!")
@@ -130,6 +182,19 @@ apple_css = """
     box-shadow: 0 2px 12px rgba(0, 0, 0, 0.08);
 }
 
+/* Model Selector */
+.model-selector {
+    margin-bottom: 24px;
+}
+
+.model-selector label {
+    font-size: 15px !important;
+    font-weight: 500 !important;
+    color: #1d1d1f !important;
+    margin-bottom: 8px !important;
+    display: block !important;
+}
+
 /* Textbox */
 textarea {
     font-size: 17px !important;
@@ -231,6 +296,10 @@ button.primary:active {
     color: #86868b !important;
 }
 
+.dark .model-selector label {
+    color: #f5f5f7 !important;
+}
+
 /* Responsive */
 @media (max-width: 734px) {
     .main-title {
@@ -276,6 +345,15 @@ with gr.Blocks(
     
     # Input Section
     with gr.Column(elem_classes="input-section"):
+        # Model Selector
+        with gr.Group(elem_classes="model-selector"):
+            model_choice = gr.Radio(
+                choices=["Base Model", "Classic Painting LoRA"],
+                value="Base Model",
+                label="Select Model",
+                info="Choose between the base Z-Image Turbo model or the Classic Painting LoRA variant"
+            )
+        
         prompt = gr.Textbox(
             placeholder="Describe the image you want to create...",
             lines=3,
@@ -306,20 +384,21 @@ with gr.Blocks(
     gr.HTML("""
         <div class="footer-text">
             <p>Powered by Z-Image Turbo from Tongyi-MAI</p>
+            <p>Classic Painting LoRA by renderartist</p>
         </div>
     """)
     
     # Event handlers
     generate_btn.click(
         fn=generate_image,
-        inputs=prompt,
+        inputs=[prompt, model_choice],
         outputs=output_image,
         api_visibility="public"
     )
     
     prompt.submit(
         fn=generate_image,
-        inputs=prompt,
+        inputs=[prompt, model_choice],
         outputs=output_image,
         api_visibility="public"
     )

requirements.txt

@@ -1,12 +1,12 @@
 spaces
-gradio
-torch
 git+https://github.com/huggingface/diffusers
 git+https://github.com/huggingface/transformers
 sentencepiece
 accelerate
 tokenizers
+gradio
 requests
 Pillow
+torch
 torchvision
 torchaudio