Update app.py

app.py

@@ -689,66 +689,66 @@ def generate_map(location_names):
     map_html = m._repr_html_()
     return map_html
 
-# device = "cuda:0" if torch.cuda.is_available() else "cpu"
-# pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16)
-# pipe.to(device)
-
-# def generate_image(prompt):
-#     with torch.cuda.amp.autocast():
-#         image = pipe(
-#             prompt,
-#             num_inference_steps=28,
-#             guidance_scale=3.0,
-#         ).images[0]
-#     return image
-
-# hardcoded_prompt_1 = "A high quality cinematic image for Toyota Truck in Birmingham skyline shot in th style of Michael Mann"
-# hardcoded_prompt_2 = "A high quality cinematic image for Alabama Quarterback close up emotional shot in th style of Michael Mann"
-# hardcoded_prompt_3 = "A high quality cinematic image for Taylor Swift concert in Birmingham skyline style of Michael Mann"
-
-# def update_images():
-#     image_1 = generate_image(hardcoded_prompt_1)
-#     image_2 = generate_image(hardcoded_prompt_2)
-#     image_3 = generate_image(hardcoded_prompt_3)
-#     return image_1, image_2, image_3
-
+#Flux Code
 
+import spaces
+import gradio as gr
+import torch
+from PIL import Image
+from diffusers import DiffusionPipeline
+import random
 
+# Initialize the base model and specific LoRA
+base_model = "black-forest-labs/FLUX.1-dev"
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 
+lora_repo = "XLabs-AI/flux-RealismLora"
+trigger_word = ""  # Leave trigger_word blank if not used.
+pipe.load_lora_weights(lora_repo)
 
+pipe.to("cuda")
 
-import torch
-from diffusers import FluxPipeline
-import os
+MAX_SEED = 2**32-1
 
-# Set PYTORCH_CUDA_ALLOC_CONF to handle memory fragmentation
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+# Hardcoded prompts for generating images
+hardcoded_prompt_1 = "A high quality cinematic image for Toyota Truck in Birmingham skyline shot in the style of Michael Mann"
+hardcoded_prompt_2 = "A high quality cinematic image for Alabama Quarterback close up emotional shot in the style of Michael Mann"
+hardcoded_prompt_3 = "A high quality cinematic image for Taylor Swift concert in Birmingham skyline style of Michael Mann"
 
-# Check GPU memory and fallback to CPU if necessary
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+@spaces.GPU(duration=80)
+def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+    # Set random seed for reproducibility
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device="cuda").manual_seed(seed)
 
+    # Update progress bar (0% saat mulai)
+    progress(0, "Starting image generation...")
 
+    # Generate image with progress updates
+    for i in range(1, steps + 1):
+        # Simulate the processing step (in a real scenario, you would integrate this with your image generation process)
+        if i % (steps // 10) == 0:  # Update every 10% of the steps
+            progress(i / steps * 100, f"Processing step {i} of {steps}...")
 
-# Reduce the inference steps and image dimensions
-def generate_image_flux(prompt, width=400, height=400, num_inference_steps=4):
+    # Generate image using the pipeline
     image = pipe(
-        prompt=prompt,
+        prompt=f"{prompt} {trigger_word}",
+        num_inference_steps=steps,
+        guidance_scale=cfg_scale,
         width=width,
         height=height,
-        num_inference_steps=num_inference_steps,  # Reduce steps to save memory
-        generator=torch.Generator("cpu").manual_seed(0),
-        guidance_scale=0.0
+        generator=generator,
+        joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
-    return image
+
+    # Final update (100%)
+    progress(100, "Completed!")
+
+    yield image, seed
 
 
-# Existing image generation function, updated to use the new Flux pipeline
 
-# Hardcoded prompts for generating images
-hardcoded_prompt_1 = "A high quality cinematic image for Toyota Truck in Birmingham skyline shot in the style of Michael Mann"
-hardcoded_prompt_2 = "A high quality cinematic image for Alabama Quarterback close up emotional shot in the style of Michael Mann"
-hardcoded_prompt_3 = "A high quality cinematic image for Taylor Swift concert in Birmingham skyline style of Michael Mann"
     
 
 
@@ -1496,12 +1496,11 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
 
         with gr.Column():
 
-            # Display generated images
-            image_output_1 = gr.Image(value=generate_image_flux(hardcoded_prompt_1), width=400, height=400)
-            image_output_2 = gr.Image(value=generate_image_flux(hardcoded_prompt_2), width=400, height=400)
-            image_output_3 = gr.Image(value=generate_image_flux(hardcoded_prompt_3), width=400, height=400)
+           # Display images generated using Flux with LoRA integration
+            image_output_1 = gr.Image(value=next(run_lora(hardcoded_prompt_1, 7.5, 50, True, 42, 512, 512, 0.5)), width=400, height=400)
+            image_output_2 = gr.Image(value=next(run_lora(hardcoded_prompt_2, 7.5, 50, True, 42, 512, 512, 0.5)), width=400, height=400)
+            image_output_3 = gr.Image(value=next(run_lora(hardcoded_prompt_3, 7.5, 50, True, 42, 512, 512, 0.5)), width=400, height=400)
 
-            # Refresh button to update images
             refresh_button = gr.Button("Refresh Images")
             refresh_button.click(fn=update_images, inputs=None, outputs=[image_output_1, image_output_2, image_output_3])