updates

app.py

@@ -1,6 +1,44 @@
+run_api = False
+SSD_1B = False
 
-#!/usr/bin/env python
+import os
+
+# Use GPU
+gpu_info = os.popen("nvidia-smi").read()
+if "failed" in gpu_info:
+    print("Not connected to a GPU")
+    is_gpu = False
+else:
+    print(gpu_info)
+    is_gpu = True
+print(is_gpu)
+
+from IPython.display import clear_output
+
+
+def check_enviroment():
+    try:
+        import torch
+
+        print("Enviroment is already installed.")
+    except ImportError:
+        print("Enviroment not found. Installing...")
+        # Install requirements from requirements.txt
+        os.system("pip install -r requirements.txt")
+        # Install gradio version 3.48.0
+        os.system("pip install gradio==3.39.0")
+        # Install python-dotenv
+        os.system("pip install python-dotenv")
+        # Clear the output
+        clear_output()
+
+        print("Enviroment installed successfully.")
 
+
+# Call the function to check and install Packages if necessary
+check_enviroment()
+
+from IPython.display import clear_output
 import os
 import gradio as gr
 import numpy as np
@@ -8,122 +46,151 @@ import PIL
 import base64
 import io
 import torch
-
-# SSD-1B
-#from diffusers import LCMScheduler, AutoPipelineForText2Image
+from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
 
 # SDXL
 from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
 
+# Get the current directory
+current_dir = os.getcwd()
+model_path = os.path.join(current_dir)
+# Set the cache path
+cache_path = os.path.join(current_dir, "cache")
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = int(os.getenv('MAX_IMAGE_SIZE', '1024'))
-SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
+SECRET_TOKEN = os.getenv("SECRET_TOKEN", "default_secret")
+
+# Uncomment the following line if you are using PyTorch 1.10 or later
+# os.environ["TORCH_USE_CUDA_DSA"] = "1"
+
+if is_gpu:
+    # Uncomment the following line if you want to enable CUDA launch blocking
+    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+else:
+    # Uncomment the following line if you want to use CPU instead of GPU
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-#device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 if torch.cuda.is_available():
+    # Get the current directory
+    current_dir = os.getcwd()
+    model_path = os.path.join(current_dir)
 
-    #pipe = AutoPipelineForText2Image.from_pretrained("segmind/SSD-1B", torch_dtype=torch.float16, variant="fp16")
-    #pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-    #pipe.to("cuda")
+    # Set the cache path
+    cache_path = os.path.join(current_dir, "cache")
 
-    # load and fuse
-    #pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
-    #pipe.fuse_lora()
+    if not SSD_1B:
 
-    unet = UNet2DConditionModel.from_pretrained("latent-consistency/lcm-sdxl", torch_dtype=torch.float16, variant="fp16")
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", unet=unet, torch_dtype=torch.float16, variant="fp16")
+        unet = UNet2DConditionModel.from_pretrained(
+            "latent-consistency/lcm-sdxl",
+            torch_dtype=torch.float16,
+            variant="fp16",
+            cache_dir=cache_path,
+        )
+        pipe = DiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            unet=unet,
+            torch_dtype=torch.float16,
+            variant="fp16",
+            cache_dir=cache_path,
+        )
+
+        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+        pipe.to("cuda")
+    else:
+        # SSD-1B
+        from diffusers import LCMScheduler, AutoPipelineForText2Image
+
+        pipe = AutoPipelineForText2Image.from_pretrained(
+            "segmind/SSD-1B",
+            torch_dtype=torch.float16,
+            variant="fp16",
+            cache_dir=cache_path,
+        )
+        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+        pipe.to("cuda")
+
+        # load and fuse
+        pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
+        pipe.fuse_lora()
 
-    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-    pipe.to('cuda')
-    
 else:
     pipe = None
 
-def generate(prompt: str,
-             negative_prompt: str = '',
-             seed: int = 0,
-             width: int = 1024,
-             height: int = 1024,
-             guidance_scale: float = 0.0,
-             num_inference_steps: int = 4,
-             secret_token: str = '') -> PIL.Image.Image:
+
+def generate(
+    prompt: str,
+    negative_prompt: str = "",
+    seed: int = 0,
+    width: int = 1024,
+    height: int = 1024,
+    guidance_scale: float = 0.0,
+    num_inference_steps: int = 4,
+    secret_token: str = "",
+) -> PIL.Image.Image:
     if secret_token != SECRET_TOKEN:
         raise gr.Error(
-            f'Invalid secret token. Please fork the original space if you want to use it for yourself.')
-        
+            f"Invalid secret token. Please fork the original space if you want to use it for yourself."
+        )
+
     generator = torch.Generator().manual_seed(seed)
 
-    image = pipe(prompt=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[0]
-    
+    image = pipe(
+        prompt=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[0]
     return image
 
-with gr.Blocks() as demo:
-    gr.HTML("""
-    <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
-        <div style="text-align: center; color: black;">
-            <p style="color: black;">This space is a REST API to programmatically generate images using LCM LoRA SSD-1B.</p>
-            <p style="color: black;">It is not meant to be directly used through a user interface, but using code and an access key.</p>
-        </div>
-    </div>""")
+
+clear_output()
+
+if not run_api:
     secret_token = gr.Text(
-        label='Secret Token',
+        label="Secret Token",
         max_lines=1,
-        placeholder='Enter your secret token',
+        placeholder="Enter your secret token",
     )
     prompt = gr.Text(
-        label='Prompt',
+        label="Prompt",
         show_label=False,
         max_lines=1,
-        placeholder='Enter your prompt',
+        placeholder="Enter your prompt",
         container=False,
     )
-    result = gr.Image(label='Result', show_label=False)
+    result = gr.Image(label="Result", show_label=False)
     negative_prompt = gr.Text(
-        label='Negative prompt',
+        label="Negative prompt",
         max_lines=1,
-        placeholder='Enter a negative prompt',
+        placeholder="Enter a negative prompt",
         visible=True,
     )
-    seed = gr.Slider(label='Seed',
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=0)
+    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
 
     width = gr.Slider(
-        label='Width',
+        label="Width",
         minimum=256,
         maximum=MAX_IMAGE_SIZE,
         step=32,
         value=1024,
     )
     height = gr.Slider(
-        label='Height',
+        label="Height",
         minimum=256,
         maximum=MAX_IMAGE_SIZE,
         step=32,
         value=1024,
     )
     guidance_scale = gr.Slider(
-        label='Guidance scale',
-        minimum=0,
-        maximum=2,
-        step=0.1,
-        value=0.0)
+        label="Guidance scale", minimum=0, maximum=2, step=0.1, value=0.0
+    )
     num_inference_steps = gr.Slider(
-        label='Number of inference steps',
-        minimum=1,
-        maximum=8,
-        step=1,
-        value=4)
-
+        label="Number of inference steps", minimum=1, maximum=8, step=1, value=4
+    )
     inputs = [
         prompt,
         negative_prompt,
@@ -134,11 +201,87 @@ with gr.Blocks() as demo:
         num_inference_steps,
         secret_token,
     ]
-    prompt.submit(
+    iface = gr.Interface(
         fn=generate,
         inputs=inputs,
         outputs=result,
-        api_name='run',
+        title="Image Generator",
+        description="Generate images based on prompts.",
     )
 
-demo.queue(max_size=32).launch()
+    iface.launch()
+
+
+if run_api:
+    with gr.Blocks() as demo:
+        gr.HTML(
+            """
+        <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
+            <div style="text-align: center; color: black;">
+                <p style="color: black;">This space is a REST API to programmatically generate images using LCM LoRA SSD-1B.</p>
+                <p style="color: black;">It is not meant to be directly used through a user interface, but using code and an access key.</p>
+            </div>
+        </div>"""
+        )
+        secret_token = gr.Text(
+            label="Secret Token",
+            max_lines=1,
+            placeholder="Enter your secret token",
+        )
+        prompt = gr.Text(
+            label="Prompt",
+            show_label=False,
+            max_lines=1,
+            placeholder="Enter your prompt",
+            container=False,
+        )
+        result = gr.Image(label="Result", show_label=False)
+        negative_prompt = gr.Text(
+            label="Negative prompt",
+            max_lines=1,
+            placeholder="Enter a negative prompt",
+            visible=True,
+        )
+        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+
+        width = gr.Slider(
+            label="Width",
+            minimum=256,
+            maximum=MAX_IMAGE_SIZE,
+            step=32,
+            value=1024,
+        )
+        height = gr.Slider(
+            label="Height",
+            minimum=256,
+            maximum=MAX_IMAGE_SIZE,
+            step=32,
+            value=1024,
+        )
+        guidance_scale = gr.Slider(
+            label="Guidance scale", minimum=0, maximum=2, step=0.1, value=0.0
+        )
+        num_inference_steps = gr.Slider(
+            label="Number of inference steps", minimum=1, maximum=8, step=1, value=4
+        )
+
+        inputs = [
+            prompt,
+            negative_prompt,
+            seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+            secret_token,
+        ]
+        prompt.submit(
+            fn=generate,
+            inputs=inputs,
+            outputs=result,
+            api_name="run",
+        )
+
+    # demo.queue(max_size=32).launch()
+    # Launch the Gradio app with multiple workers and debug mode enabled
+    demo.queue(max_size=32).launch(debug=True)