Update app.py

app.py

@@ -4,15 +4,11 @@ import uuid
 import gradio as gr
 import numpy as np
 from PIL import Image
+import spaces
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler, FluxPipeline
-from huggingface_hub import snapshot_download
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 from typing import Tuple
 
-# Ensure Hugging Face token from secrets
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-# Function to apply the style based on the selected model
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     styles = {
         "3840 x 2160": (
@@ -26,129 +22,99 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     return p.replace("{prompt}", positive), n + negative
 
-# Function to load and prepare the model
-def load_and_prepare_model(model_name: str):
-    if model_name == "RealVisXL_V5.0_Lightning":
-        model_id = "SG161222/RealVisXL_V5.0_Lightning"
-        # Ensure the model is downloaded locally
-        local_model_path = snapshot_download(model_id, token=HF_TOKEN)
-        pipe = StableDiffusionXLPipeline.from_pretrained(
-            local_model_path,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            use_safetensors=True,
-            add_watermarker=False,
-        ).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    elif model_name == "FLUX.1-dev":
-        model_id = "black-forest-labs/FLUX.1-dev"
-        # Ensure the model is downloaded locally
-        local_model_path = snapshot_download(model_id, token=HF_TOKEN)
-        pipe = FluxPipeline.from_pretrained(
-            local_model_path,
-            torch_dtype=torch.bfloat16,
-        )
-        pipe.enable_model_cpu_offload()  # Save VRAM by offloading model to CPU
-    else:
-        raise ValueError("Unsupported model")
-    
+def load_and_prepare_model():
+    model_id = "SG161222/RealVisXL_V5.0_Lightning"
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        use_safetensors=True,
+        add_watermarker=False,
+    ).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     return pipe
 
-# Function to save image
+model = load_and_prepare_model()
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, np.iinfo(np.int32).max)
+    return seed
+
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
 
-# Main image generation function
+@spaces.GPU(duration=60, enable_queue=True)
 def generate(
         prompt: str,
-        model_name: str,
         seed: int = 1,
         width: int = 1024,
         height: int = 1024,
         guidance_scale: float = 3,
-        num_inference_steps: int = 50,
+        num_inference_steps: int = 25,
         randomize_seed: bool = False,
 ):
-    model = load_and_prepare_model(model_name)
-    seed = random.randint(0, np.iinfo(np.int32).max) if randomize_seed else seed
-    generator = torch.Generator("cpu" if model_name == "FLUX.1-dev" else model.device).manual_seed(seed)
+    global model
 
-    if model_name == "RealVisXL_V5.0_Lightning":
-        positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt)
-        options = {
-            "prompt": [positive_prompt],
-            "negative_prompt": [negative_prompt],
-            "width": width,
-            "height": height,
-            "guidance_scale": guidance_scale,
-            "num_inference_steps": num_inference_steps,
-            "generator": generator,
-            "output_type": "pil",
-        }
-        images = model(**options).images
-    elif model_name == "FLUX.1-dev":
-        image = model(
-            prompt=prompt,
-            height=height,
-            width=width,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            max_sequence_length=512,
-            generator=generator
-        ).images[0]
-        images = [image]
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator(device=model.device).manual_seed(seed)
 
-    image_paths = [save_image(img) for img in images]
-    return image_paths, seed
+    positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt)
+
+    options = {
+        "prompt": [positive_prompt],
+        "negative_prompt": [negative_prompt],
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "generator": generator,
+        "output_type": "pil",
+    }
+
+    images = model(**options).images
+    image_path = save_image(images[0])  # Saving the first generated image
+    return image_path
 
-# Gradio interface setup
 with gr.Blocks(theme="soft") as demo:
     # Centered text "SNAPSCRIBE" at the top of the screen
     gr.Markdown("<h1 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>SNAPSCRIBE</h1>")
     
-    # Dropdown for model selection
     with gr.Row():
         with gr.Column(scale=3):
-            model_dropdown = gr.Dropdown(
-                choices=["RealVisXL_V5.0_Lightning", "FLUX.1-dev"],
-                label="Select Model",
-                value="RealVisXL_V5.0_Lightning"
-            )
             prompt = gr.Textbox(
                 label="Input Prompt",
                 placeholder="Describe the image you want to create",
                 lines=2,
             )
             run_button = gr.Button("Generate Image")
+            gr.Markdown("Developed using the RealVisXL_V5.0_Lightning model.", elem_id="model_info")
         with gr.Column(scale=7):
-            result = gr.Gallery(label="Generated Image", columns=2)
+            result_image = gr.Image(label="Generated Image", type="filepath")
 
     run_button.click(
         fn=generate,
-        inputs=[prompt, model_dropdown],
-        outputs=[result],
+        inputs=[prompt],
+        outputs=[result_image],
     )
 
-    # Footer added to center-align the text
-    gr.HTML("""
+    # Footer with custom style and text
+    gr.Markdown("""
     <style>
-        .footer {
-            position: relative;
-            left: 0;
-            bottom: 0;
-            width: 100%;
-            background-color: white;
-            color: black;
-            text-align: center;
-            padding: 10px;
-            margin-top: 20px;
-        }
+    .footer {
+        position: relative;
+        left: 0;
+        bottom: 0;
+        width: 100%;
+        background-color: white;
+        color: black;
+        text-align: center;
+    }
     </style>
     <div class="footer">
-        <p>Developed with ❤ by Aklavya(Bucky)</p>
+    <p>Developed with ❤ by Aklavya(Bucky)</p>
     </div>
     """)
 
-# Launch the Gradio interface
 demo.launch()