Fix UI

app.py

@@ -6,6 +6,7 @@ from diffusers import StableDiffusionDepth2ImgPipeline
 from pathlib import Path
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+is_gpu_associated = torch.cuda.is_available()
 dept2img = StableDiffusionDepth2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-2-depth",
     torch_dtype=torch.float16,
@@ -38,6 +39,9 @@ def pad_image(input_image):
 
 
 def predict(input_image, prompt, negative_prompt, steps, num_samples, scale, seed, strength, depth_image=None):
+    if not is_gpu_associated:
+        raise gr.Error("Please associate a T4 GPU for this Space")
+    torch.cuda.empty_cache()
     depth = None
     if depth_image is not None:
         depth_image = pad_image(depth_image)
@@ -66,20 +70,48 @@ def predict(input_image, prompt, negative_prompt, steps, num_samples, scale, see
 block = gr.Blocks().queue()
 with block:
     with gr.Box():
-        top_description = gr.HTML(f'''
-                <div style="text-align: center; max-width: 650px; margin: 0 auto;">
-                <div>
-                    <img class="logo" src="file/mirage.png" alt="Mirage Logo"
-                        style="margin: auto; max-width: 7rem;">
+        if is_gpu_associated:
+                top_description = gr.HTML(f'''
+                    <div style="text-align: center; max-width: 650px; margin: 0 auto;">
+                    <div>
+                        <img class="logo" src="file/mirage.png" alt="Mirage Logo"
+                            style="margin: auto; max-width: 7rem;">
+                        <br />
+                        <h1 style="font-weight: 900; font-size: 2.5rem;">
+                        Depth2Img Web UI
+                        </h1>
+                        <br />
+                        <a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/MirageML/depth2img?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a>
+                    </div>
                     <br />
-                    <h1 style="font-weight: 900; font-size: 2.5rem;">
-                    Depth2Img Web UI
-                    </h1>
                     <br />
-                    <a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/MirageML/dreambooth?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a>
-                </div>
-                </div>
-            ''')
+                    <p style="margin-bottom: 10px; font-size: 94%">
+                    Create variations of an image while preserving shape and depth!
+                    </p>
+                    </div>
+                ''')
+            else:
+                top_description = gr.HTML(f'''
+                    <div style="text-align: center; max-width: 650px; margin: 0 auto;">
+                    <div>
+                        <img class="logo" src="file/mirage.png" alt="Mirage Logo"
+                            style="margin: auto; max-width: 7rem;">
+                        <br />
+                        <h1 style="font-weight: 900; font-size: 2.5rem;">
+                        Depth2Img Web UI
+                        </h1>
+                        <br />
+                        <a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/MirageML/depth2img?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a>
+                    </div>
+                    <br />
+                    <br />
+                    <p style="margin-bottom: 10px; font-size: 94%">
+                    Create variations of an image while preserving shape and depth!
+                    </p>
+                    <br />
+                    <p>There's only one step left before you can run the app: <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings" style="text-decoration: underline" target="_blank">attribute a <b>T4 GPU</b> to it (via the Settings tab)</a> and run the training below. You will be billed by the minute from when you activate the GPU until when it is turned it off.</p>
+                    </div>
+                ''')
 
     with gr.Row():
         with gr.Column():
@@ -112,7 +144,7 @@ with block:
         with gr.Column():
             gallery = gr.Gallery(label="Generated images", show_label=False).style(
                 grid=[2], height="auto")
-    if device == "cuda":
+    if is_gpu_associated:
         gr.Examples(
             examples=[
                 ["./examples/original_iso.png", "hogwarts castle",