Update app.py

app.py

@@ -1,14 +1,15 @@
 import spaces
 import gradio as gr
 import torch
-from diffusers import ZImagePipeline, AutoPipelineForText2Image
+from diffusers import ZImagePipeline
 import os
 from pathlib import Path
 
-# Load the base model directly at startup
+# Load the 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",
@@ -23,27 +24,21 @@ 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 with optional LoRA.
+    Generate an image using Z-Image Turbo model.
     
     Args:
         prompt: Text description of the desired image
-        model_choice: Either "Base Model" or "Classic Painting LoRA"
     
     Returns:
         Generated PIL Image
     """
-    global pipe, current_model, lora_loaded
+    global pipe
     
     if pipe is None:
         raise gr.Error("Model failed to load on startup. Please restart the application.")
@@ -54,69 +49,22 @@ def generate_image(
     # Determine device
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
-    # Handle model switching
-    progress(0.05, desc="Loading model...")
+    # Set random seed for reproducibility
+    generator = torch.Generator(device).manual_seed(42)
+    
+    # Generate the image with optimal settings
+    progress(0.1, desc="Generating image...")
     
     try:
-        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,
-            )
+        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!")
@@ -182,19 +130,6 @@ 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;
@@ -296,10 +231,6 @@ button.primary:active {
     color: #86868b !important;
 }
 
-.dark .model-selector label {
-    color: #f5f5f7 !important;
-}
-
 /* Responsive */
 @media (max-width: 734px) {
     .main-title {
@@ -345,15 +276,6 @@ 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,
@@ -384,21 +306,20 @@ 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, model_choice],
+        inputs=prompt,
         outputs=output_image,
         api_visibility="public"
     )
     
     prompt.submit(
         fn=generate_image,
-        inputs=[prompt, model_choice],
+        inputs=prompt,
         outputs=output_image,
         api_visibility="public"
     )