ok

main.py

@@ -67,6 +67,7 @@ class GenerationRequest(BaseModel):
     temp: float = 1.0
     top_k: int = 256
     grid_size: int = 4
+    advanced: bool = False
 
 class GenerationOutput(BaseModel):
     output: List[str]
@@ -85,7 +86,11 @@ def generate(request: GenerationRequest):
         else:
             return ErrorOutput(error=output[1])
         
-        images = model.sample(request)
+        if request.advanced:
+            images = model.sample_two_stages(request)
+        else:
+            images = model.sample(request)
+            
         images_encoded = []
         for x in range(len(images)):
             image = simplejpeg.encode_jpeg(images[x], quality=95)