Update app.py

app.py

@@ -1,146 +1,101 @@
 import gradio as gr
-import numpy as np
-import random
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
-    
-    image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
-    
-    return image
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 520px;
-}
-"""
-
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
-with gr.Blocks(css=css) as demo:
+from PIL import Image, ImageColor
+from back_task import *
+
+# The function that does the hard work
+def generate(radio,color,prompt, guidance_loss_scale):
+    print(color)
+    if radio == "color guidance":
+        target_color = ImageColor.getcolor(color, "RGB")  # Target color as RGB
+        target_color = [a / 255 for a in target_color]  # Rescale from (0, 255) to (0, 1)
+
+    elif radio == "text guidance":
+        # We embed a prompt with CLIP as our target
+        text = open_clip.tokenize([prompt]).to(device)
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            text_features = clip_model.encode_text(text)
+
+    x = torch.randn(1, 3, 256, 256).to(device)
+    for i, t in tqdm(enumerate(scheduler.timesteps)):
+        model_input = scheduler.scale_model_input(x, t)
+        with torch.no_grad():
+            noise_pred = image_pipe.unet(model_input, t)["sample"]
+
+        if radio == "color guidance":
+            x = x.detach().requires_grad_()
+            x0 = scheduler.step(noise_pred, t, x).pred_original_sample
+            loss = color_loss(x0, target_color) * guidance_loss_scale
+            cond_grad = -torch.autograd.grad(loss, x)[0]
+            x = x.detach() + cond_grad
+        elif radio == "text guidance":
+                cond_grad = 0
+
+                for cut in range(n_cuts):
+
+                    # Set requires grad on x
+                    x = x.detach().requires_grad_()
+
+                    # Get the predicted x0:
+                    x0 = scheduler.step(noise_pred, t, x).pred_original_sample
+
+                    # Calculate loss
+                    loss = clip_loss(x0, text_features) * guidance_loss_scale
+
+                    # Get gradient (scale by n_cuts since we want the average)
+                    cond_grad -= torch.autograd.grad(loss, x)[0] / n_cuts
+
+
+                # Modify x based on this gradient
+                alpha_bar = scheduler.alphas_cumprod[i]
+                x = x.detach() + cond_grad * alpha_bar.sqrt()  # Note the additional scaling factor here!
+
     
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""
-        # Text-to-Image Gradio Template
-        Currently running on {power_device}.
-        """)
-        
-        with gr.Row():
-            
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            
-            run_button = gr.Button("Run", scale=0)
-        
-        result = gr.Image(label="Result", show_label=False)
+        x = scheduler.step(noise_pred, t, x).prev_sample
+    grid = torchvision.utils.make_grid(x, nrow=4)
+    im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
+    im = Image.fromarray(np.array(im * 255).astype(np.uint8))
+    # im.save("test.jpeg")
+    return im
 
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-            
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
-            with gr.Row():
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
+
+title="""<h1 align="center">Make me a WikiArt</h1>
+<p align="center">Try-out of exercise from HF Learn [Difussion Course] </p>
+<p><center>
+<a href="https://huggingface.co/learn/diffusion-course" target="_blank">[HF-Learn]</a>
+</center></p>"""
+
+with gr.Blocks() as demo:
+    gr.HTML(title)
+    with gr.Row():
+        with gr.Column():
+            # Create a radio button with options "no guidance", "color guidance", and "text guidance"
+            radio = gr.Radio(["no guidance", "color guidance", "text guidance"], label="Choose",value="no guidance")
+
+            # Create a textbox that only shows when 'text guidance' is selected
+            text = gr.Textbox(label="This text only shows when 'text guidance' is selected.", visible=False)
+
+            # Create a color picker (not a tuple)
+            color = gr.ColorPicker(label="color", value="#000000", visible=False)
+
+            # Create a slider that shows when any option is selected
+            slider = gr.Slider(label="guidance_scale", minimum=0, maximum=30, value=3, visible=False)
+
+            def update_visibility(radio):
+                value = radio  # Get the selected value from the radio button
+                if value == "color guidance":
+                    return [gr.Textbox(visible=False),gr.ColorPicker(visible=True),gr.Slider(visible=True)] #make it visible
+                elif value == "text guidance":
+                    return [gr.Textbox(visible=True),gr.ColorPicker(visible=False),gr.Slider(visible=True)]
+                else:
+                    return [gr.Textbox(visible=False),gr.ColorPicker(visible=False),gr.Slider(visible=False)]
+
+            radio.change(update_visibility, radio,[text,color,slider])
+        with gr.Column():
+            outputs = gr.Image(label="result")
             
-            with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,
-                )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=12,
-                    step=1,
-                    value=2,
-                )
-        
-        gr.Examples(
-            examples = examples,
-            inputs = [prompt]
-        )
-
-    run_button.click(
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
-    )
+    with gr.Row():
+        gen_bttn=gr.Button(value="generate")
+        gen_bttn.click(generate, inputs=[radio,color,text,slider], outputs=outputs)
+
 
 demo.queue().launch()